llama.cppにおけるプロンプト検索ドラフトの高速化

llama.cppにおけるプロンプト検索ドラフトの42倍高速化 Hayder Tirmazi [ホームページ] [github] [twitter] この記事は2026年9月26日に公開されました。 TL;DR Daniel Lemire氏とMartin Ankerl氏の研究に基づいた一連のシンプルなパフォーマンス最適化により、llama.cppにおけるプロンプト検索デコーディングのドラフト作成を最大42倍高速化し、メモリ使用量を最大2.6倍削減しました。更新: Daniel Lemire氏がプルリクエストを送り、元の最適化に加えてプロンプト検索ドラフトを最大4.2倍高速化しました。これにより、全体の速度向上は最大140倍になります。詳細は以下で説明します。 多くの人気のある推論エンジン（llama.cppやvllmなど）や、Hugging FaceのTransformersライブラリのような機械学習ライブラリは、トークン生成を高速化するためにプロンプト検索デコーディング（n-gramスペキュレーションとも呼ばれる）をサポートしています。プロンプト検索デコーディングは、技術的にはスペキュラティブデコーディングの特殊なケースであり、非常に単純なドラフトモデル、すなわちn-gramモデルを使用します。 プロンプト検索デコーディングが使用される際、推論エンジンは以下のルールに従って次の$k$個のトークンをドラフトします。モデルの現在のトークンを$x_1, ext{ldots}, x_t$とします。n-gramは連続する$n$個のトークンのシーケンスです。例えば、3-gramは$(x_1, x_2, x_3)$または$(x_2, x_3, x_4)$、一般的には任意の$i ext{ in } ext{{1, ext{ldots}, t-2}}$に対して$(x_i, x_{i+1}, x_{i+2})$となります。n-gramモデルは、前の$n-1$個のトークンに基づいて次のトークンを予測する確率モデルです。アイデアは非常にシンプルです。まず、テキストのコーパスを選択し、それをn-gramに解析します。次に、各n-gramの頻度をカウントします。n-gramモデルが$n-1$個のトークンのシーケンスの後に次のトークンを予測する必要がある場合、コーパス内でそのシーケンスの後に最も頻繁に出現するトークンを選択させます。 lama.cppは3種類のn-gramキャッシュを維持しています。任意のn-gramを$ ext{eta}$、任意のトークンを$y$とします。n-gramキャッシュは、与えられたコーパスと語彙において、すべてのn-gram $ ext{eta}$とすべてのトークン $y$に対して、トークン $y$がn-gram $ ext{eta}$の後に何回出現するかを示す$c( ext{eta}, y)$を格納するデータ構造です。llama.cppが使用する3つのn-gramキャッシュは、コンテキストキャッシュ、ダイナミックキャッシュ、スタティックキャッシュです。 コンテキストキャッシュは、モデルによって処理されている現在のトークン$x_1, ext{ldots}, x_t$のサイズ1から4までのn-gramを格納します。コンテキストキャッシュは、モデルが新しいトークンを生成するにつれて更新されます。ダイナミックキャッシュは、モデルの以前の実行（例えば、以前の会話）からのn-gramのカウントを格納します。最後に、スタティックキャッシュは、llama-lookup-createで構築された静的テキストコーパスからのサイズ2のn-gramを格納します。コンテキスト、ダイナミック、スタティックキャッシュをそれぞれ$c_{ ext{ctx}}$、$c_{ ext{dyn}}$、$c_{ ext{st}}$とします。 lama.cppは、n-gramキャッシュを使用して次のように新しいトークンをドラフトします。モデルによって処理された前の$n$個のトークンを$X_n = (x_{t-n+1}, ext{ldots}, x_t)$とします。語彙内のすべてのトークン $y$に対して、llama.cppは以下の式を使用してスコアを計算します。 $$s_n^{f}(y) = f(X_n, y) ext{cdot} w(y) ext{quad} ext{where} ext{quad} w(y) = egin{cases} 100 ext{ } c_{ ext{st}}(X_2, y) & ext{if } c_{ ext{st}}(X_2, y) > 0 \ 1 & ext{otherwise} ext{endcases}$$ ここで、$f$はコンテキストキャッシュ$c_{ ext{ctx}}$またはダイナミックキャッシュ$c_{ ext{dyn}}$のいずれかです。重み$w(y)$は、スタティックキャッシュとも一致するトークンを優先することに注意してください。スタティックキャッシュがない場合、$w(y) = 1$となります。 各$n$に対して、llama.cppは最もスコアの高いトークン$y^* = ext{argmax}_y s_n^{f}(y)$を選択します。$F(X_n) = ext{sum}_y f(X_n, y)$を、$X_n$がトークンと共に現れた回数とします。llama.cppは、2つの設定可能な閾値$a_n$と$p_n$に基づいて、次のように$y^*$をドラフトします。 $$F(X_n) ext{ge} a_n ext{quad} ext{and} ext{quad} f(X_n, y^*) ext{ge} p_n ext{ } F(X_n)$$ 言い換えれば、$X_n$が少なくとも$a_n$回出現し、かつトークン$y^*$がそれらの出現の少なくとも$p_n$の割合で$X_n$に続いた場合に、$y^*$はドラフトトークンとして受け入れられます。 lama.cppのリリースb11182の時点では、閾値は次のようにハードコードされています。コンテキストキャッシュの場合、$(a_1, a_2, a_3, a_4) = (2, 2, 1, 1)$および$(p_1, p_2, p_3, p_4) = (0.66, 0.5, 0.5, 0.5)$です。ダイナミックキャッシュの場合、$(a_1, a_2, a_3, a_4) = (4, 3, 2, 2)$および$(p_1, p_2, p_3, p_4) = (0.75, 0.66, 0.66, 0.66)$です。 lama.cppは$n = 4, 3, 2, 1$を試行し、上記の条件を最初に満たす$y^*$をドラフトします。まず$c_{ ext{ctx}}$でスコアリングします。$c_{ ext{ctx}}$からの候補がどの$n$でもパスしない場合にのみ$c_{ ext{dyn}}$でスコアリングします。$c_{ ext{dyn}}$からの候補もパスしない場合、llama.cppはスタティックキャッシュのみに依存するようになります（他のキャッシュの候補を再重み付けするためだけに使用するのではなく）。 補足ですが、llama.cppにおけるスタティックキャッシュの閾値は、対応する$n$（すなわち$n=2$）に対するコンテキストキャッシュの閾値と同じ値です。$C_{ ext{st}}(X_2) = ext{sum}_y c_{ ext{st}}(X_2, y)$とします。llama.cppは、$C_{ ext{st}}(X_2) ext{ge} a_2 = 2$および$c_{ ext{st}}(X_2, y) ext{ge} p_2 ext{ } C_{ ext{st}}(X_2) = 0.5 ext{ } C_{ ext{st}}(X_2)$の場合に、$c_{ ext{st}}(X_2, y)$が最大のトークン$y$を選択し、それをドラフトします。スタティックキャッシュも失敗した場合、llama.cppは次のトークンをドラフトしません。 実験設定 lama.cppのリポジトリには、プロンプト検索デコーディングの例が含まれています。ここでは、コーパスからスタティックキャッシュを構築するためのllama-lookup-createと、プロンプト検索デコーディングのベンチマークを行うためのllama-lookup-statsという2つのツールを使用します。llama-lookup-statsは基本的にファイルを読み込み、ファイルのトークンをモデルの出力として扱います。llama.cppのドラフティングループをシミュレートされた「モデル出力」（つまりファイル）上で実行し、ドラフトされたトークンのうちいくつがファイルと一致したか、トークンをドラフトするのにかかった時間、スタティックngramキャッシュのロードにかかった時間を記録します。 WikiText-103を使用してスタティックキャッシュを構築し、WikiText-103のテストテキストをllama-lookup-statsでリプレイします。この評価方法は、llama.cppにスタティックn-gramキャッシュを追加した@JohannesGaessler氏のPRから借用しました。私の変更はプロンプト検索デコーディングのアルゴリズムに影響を与えないため、データセットは主に受け入れ率に影響しますが、私の変更は受け入れ率を変更しません。念のため、私の変更が元の実装とほぼ同じ受け入れ率を維持していることを確認します。ここで実際に変更される重要なメトリックは、1) ドラフトトークンあたりのレイテンシ、2) スタティックキャッシュのロード時間、3) スタティックキャッシュのメモリ使用量です。 スタティックn-gramキャッシュの異なるコーパスサイズでパフォーマンスがどのように変化するかを観察したいとも思いました。そのため、約541MBのWikiText-103の全コーパスを評価するのに加えて、WikiText-103のトレーニングテキストの最初の25MB、50MB、100MB、200MBからもスタティックキャッシュを構築しました。図中のコーパスサイズ0は、スタティックキャッシュなしで実行した場合を示しており、コンテキストキャッシュとダイナミックキャッシュのみを測定します。 この作業のすべての結果について、3回の実行の中央値を報告し、エラーバーは実行の最小値と最大値を示します。前述のllama.cppのPRに従い、モデルのコンテキストサイドを4096トークンと仮定してベンチマークを実行します。すべての実験は、14コア、48GBメモリのApple M4 Proで実行しました。私のコードと結果はすべてこのリポジトリにあります。 マップのコピー停止 lama.cppのn-gramキャッシュは現在、ネストされたstd::unordered_mapsとして実装されています。外側のマップは、各n-gramを、それに続くトークンとそのカウントのマップにマッピングします。これは、最適化というよりはバグ修正に近いものです。ドラフティングステップの複数の場所で、内側のマップが不必要にコピーされていることがわかりました。これを参照で読み取るように簡単なPRを作成しました。これにより、コーパスのサイズに応じて、ドラフティングが4.5倍から25.6倍高速になりました（下の図を参照）。レイテンシは、ドラフトされたトークンあたりの平均ドラフト時間です。 外側のマップ -> フラットハッシュマップ lama.cppは、マップのマップとしてn-gramキャッシュを実装しています。 typedef std::unordered_map<common_ngram, common_ngram_cache_part