UniPrefill: ブロック単位の動的スパース化による汎用的長コンテキストプリフィル高速化

要旨

大規模言語モデル（LLM）が急速に進歩し続けるにつれて、その能力は向上すると同時に、ますます長いコンテキスト長が要求されるようになっている。長いコンテキスト処理の推論効率を向上させるために、最近いくつかの新しい低複雑性ハイブリッドアーキテクチャが提案され、長いコンテキスト推論の計算負荷を効果的に軽減している。しかし、長いコンテキストのプリフィル高速化に関する既存の研究は、主にスパースアテンション機構に焦点を当てており、これらはフルアテンションモデルでのみ最大の高速化を達成する。これらのプリフィル高速化手法は、リニア／フルアテンションハイブリッドやスライディングウィンドウ／フルアテンションハイブリッドといった新しいアーキテクチャに適用されると、性能が大幅に低下する。さらに、そのような手法は一般に連続バッチ処理と互換性がなく、vLLMのような現代の推論エンジンへの統合が困難である。そこで本稿では、ほぼすべてのモデルアーキテクチャに適用可能なプリフィル高速化フレームワークUniPrefillを提案する。これはトークンレベルでモデルの計算を直接高速化する。さらに、UniPrefillを連続バッチ演算子として実装し、vLLMのスケジューリング戦略を拡張して、プリフィルとデコードの共同処理およびテンソル並列処理をネイティブにサポートすることで、vLLMへのシームレスな統合を実現する。UniPrefillは、Time-To-First-Token（TTFT）において最大2.1倍の高速化を達成し、同時リクエスト数が増加するにつれてその高速化効果はさらに顕著になる。

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.