UniPrefill: Aceleração Universal de Preenchimento de Contexto Longo via Esparsificação Dinâmica por Blocos

Resumo

À medida que os modelos de linguagem de grande porte (LLMs) continuam a avançar rapidamente, tornam-se cada vez mais capazes, ao mesmo tempo que exigem comprimentos de contexto cada vez maiores. Para melhorar a eficiência de inferência no processamento de contextos longos, várias arquiteturas híbridas de baixa complexidade foram recentemente propostas, aliviando efetivamente a carga computacional da inferência com contextos longos. No entanto, a pesquisa existente sobre aceleração de pré-preenchimento para contextos longos permanece predominantemente focada em mecanismos de atenção esparsa, que atingem sua aceleração máxima apenas em modelos de atenção completa. Quando transferidas para arquiteturas emergentes — como híbridos de atenção linear/completa ou híbridos de atenção com janela deslizante/completa — essas abordagens de aceleração de pré-preenchimento sofrem degradação significativa de desempenho. Além disso, tais métodos são geralmente incompatíveis com o agrupamento contínuo, dificultando sua integração em motores de inferência modernos, como o vLLM. Para isso, propomos o UniPrefill, uma estrutura de aceleração de pré-preenchimento aplicável a praticamente qualquer arquitetura de modelo, que acelera diretamente o cálculo do modelo no nível do token. Implementamos ainda o UniPrefill como um operador de agrupamento contínuo e estendemos a estratégia de escalonamento do vLLM para suportar nativamente o coprocessamento pré-preenchimento/decodificação e o paralelismo de tensor para o UniPrefill, permitindo sua integração contínua no vLLM. O UniPrefill alcança até 2,1x de aceleração no Tempo até o Primeiro Token (TTFT), com a aceleração se tornando cada vez mais pronunciada à medida que o número de solicitações concorrentes cresce.

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.