UniPrefill : Accélération universelle du pré-remplissage de longs contextes via la sparsification dynamique par blocs

Résumé

Alors que les grands modèles de langage (LLMs) continuent de progresser rapidement, ils deviennent de plus en plus performants tout en exigeant des longueurs de contexte toujours plus longues. Pour améliorer l'efficacité d'inférence du traitement de contextes longs, plusieurs architectures hybrides à faible complexité ont récemment été proposées, atténuant efficacement le fardeau computationnel de l'inférence sur de longs contextes. Cependant, les recherches existantes sur l'accélération du préremplissage (prefill) pour les contextes longs restent principalement centrées sur les mécanismes d'attention sparse, qui n'atteignent leur accélération maximale que sur les modèles à attention complète. Transférées aux architectures émergentes—telles que les hybrides d'attention linéaire/attention complète ou les hybrides de fenêtre glissante/attention complète—ces approches d'accélération du préremplissage subissent une dégradation significative de leurs performances. De plus, ces méthodes sont généralement incompatibles avec le regroupement continu (continuous batching), ce qui les rend difficiles à intégrer dans des moteurs d'inférence modernes comme vLLM. À cette fin, nous proposons UniPrefill, un cadre d'accélération du préremplissage applicable à pratiquement toute architecture de modèle, qui accélère directement le calcul du modèle au niveau des jetons. Nous implémentons également UniPrefill en tant qu'opérateur de regroupement continu et étendons la stratégie d'ordonnancement de vLLM pour prendre en charge nativement le co-traitement préremplissage-décodage et le parallélisme de tenseur pour UniPrefill, permettant ainsi son intégration transparente dans vLLM. UniPrefill atteint un gain de vitesse allant jusqu'à 2,1x en temps jusqu'au premier jeton (TTFT), l'accélération devenant de plus en plus prononcée à mesure que le nombre de requêtes concurrentes augmente.

English

As large language models (LLMs) continue to advance rapidly, they are becoming increasingly capable while simultaneously demanding ever-longer context lengths. To improve the inference efficiency of long-context processing, several novel low-complexity hybrid architectures have recently been proposed, effectively alleviating the computational burden of long-context inference. However, existing research on long-context prefill acceleration remains predominantly focused on sparse attention mechanisms, which achieve their maximum speedup only on full-attention models. When transferred to emerging architectures--such as linear/full attention hybrids or sliding window/full attention hybrids--these prefill acceleration approaches suffer significant performance degradation. Furthermore, such methods are generally incompatible with continuous batching, making them difficult to integrate into modern inference engines such as vLLM. To this end, we propose UniPrefill, a prefill acceleration framework applicable to virtually any model architecture, which directly accelerates the model's computation at the token level. We further implement UniPrefill as a continuous batching operator and extend vLLM's scheduling strategy to natively support prefill-decode co-processing and tensor parallel for UniPrefill, enabling its seamless integration into vLLM. UniPrefill achieves up to 2.1x speedup in Time-To-First-Token (TTFT), with the acceleration becoming increasingly pronounced as the number of concurrent requests grows.