Prepacking: A Simple Method for Fast Prefilling and Increased Throughput in Large Language Models

要約

トランスフォーマーベースの大規模言語モデル (LLM) の推論中、事前入力とは、自己回帰生成の前にプロンプ​​ト内の入力トークンのキー/値 (KV) キャッシュを計算することです。
入力プロンプトの長さが長い場合、事前入力によりデコード時間に大幅なオーバーヘッドが発生します。
この研究では、プレフィルの次の落とし穴に焦点を当てます。大きく変化するプロンプト長を含むバッチの場合、シーケンスを最大長までパディングする標準的な方法により、かなりの計算が無駄になります。
LLM がより長いコンテキスト長 (潜在的には最大 1,000 万トークン) をサポートするようになるにつれて、バッチ内のプロンプト長の変動がより顕著になります。
これに対処するために、プレフィル計算を最適化するためのシンプルかつ効果的な方法であるプレパッキングを提案します。
パッド トークンの冗長な計算を回避するために、プリパッキングではさまざまな長さのプロンプトをシーケンスに結合し、ビン パッキング アルゴリズムを使用して複数のシーケンスをコンパクトなバッチにパックします。
次に、アテンション マスクと位置エンコーディングを変更して、単一シーケンス内の複数のプロンプトに対して事前に入力された複数の KV キャッシュを計算します。
さまざまな長さのプロンプトを含む標準の厳選されたデータセットでは、さまざまな基本モデル構成と推論提供シナリオにわたって、Huggingface 内のデフォルトのパディングベースの事前入力計算と比較して、速度とメモリ効率が大幅に向上しました。

要約(オリジナル)

During inference for transformer-based large language models (LLM), prefilling is the computation of the key-value (KV) cache for input tokens in the prompt prior to autoregressive generation. For longer input prompt lengths, prefilling will incur a significant overhead on decoding time. In this work, we highlight the following pitfall of prefilling: for batches containing high-varying prompt lengths, significant computation is wasted by the standard practice of padding sequences to the maximum length. As LLMs increasingly support longer context lengths, potentially up to 10 million tokens, variations in prompt lengths within a batch become more pronounced. To address this, we propose Prepacking, a simple yet effective method to optimize prefilling computation. To avoid redundant computation on pad tokens, prepacking combines prompts of varying lengths into a sequence and packs multiple sequences into a compact batch using a bin-packing algorithm. It then modifies the attention mask and positional encoding to compute multiple prefilled KV-caches for multiple prompts within a single sequence. On standard curated dataset containing prompts with varying lengths, we obtain a significant speed and memory efficiency improvements as compared to the default padding-based prefilling computation within Huggingface across a range of base model configurations and inference serving scenarios.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google