llama.cppのn-gramキャッシュを最大140倍高速化——「マップを参照渡しにするだけ」で25倍速になった最適化の記録

9月26日、Hayder Tirmziが「42x Faster Prompt Lookup Drafting in llama.cpp」と題した記事を公開した。llama.cppにおけるPrompt Lookup Decodingのドラフティングを最大42倍高速化し、メモリ使用量を最大2.6倍削減した手法について詳しく紹介している。

最も驚くべき事実は、4つの最適化のうち最初の1つ——std::unordered_mapの内側マップが不必要にコピーされていた実装を参照渡しに直すだけ——で、コーパスサイズに応じてドラフティングが4.5倍〜25.6倍高速化されたことだ。記事公開後には著名なパフォーマンス研究者のDaniel Lemireもさらなる最適化を投入し、アップストリームのllama.cppと比較した総合的な高速化は最大140倍に達した。

背景：なぜllama.cppの推論速度が重要なのか

llama.cppはCPUやApple Siliconでローカル実行できるLLM推論エンジンとして広く普及している。クラウドAPIへの依存を避けたいユーザーや、エッジデバイスでの推論を検討する開発者にとって、その推論速度は直接的なユーザー体験に直結する。特に長いプロンプトを扱うRAGや要約タスクでは、トークン生成のレイテンシ改善が実用性の鍵を握る。

Prompt Lookup Decodingとは何か

Prompt Lookup Decoding（n-gram投機的デコーディングとも呼ばれる）は、llama.cppやvLLM、Hugging FaceのTransformersライブラリといった主要な推論エンジンが対応するトークン生成の高速化手法だ。

投機的デコーディング（Speculative Decoding）は、軽量なドラフトモデルで次のトークン候補を先読みし、本モデルで一括検証することでスループットを向上させる手法だ。Prompt Lookup Decodingはその派生で、ドラフトモデルの代わりにn-gramマッチングを使うため、追加のモデルを必要とせず実装が単純なのが特徴だ。

llama.cppはこのために3種類のn-gramキャッシュを内部に持つ。

コンテキストキャッシュ：現在処理中のトークン列から構築（1〜4-gram）
ダイナミックキャッシュ：過去の会話履歴から蓄積
静的キャッシュ：llama-lookup-createで事前に構築した外部テキストコーパスから生成（2-gram）

Tirmziが今回最適化したのは、この静的キャッシュのドラフティング処理だ。実験環境はApple M4 Pro（14コア、48GBメモリ）、ベンチマークにはWikiText-103コーパス（最大541MB）を使用している。

最も効いた最適化：マップの不要なコピーを止める

最初の最適化が最も劇的な効果をもたらした。

llama.cppのn-gramキャッシュはstd::unordered_mapのネスト構造で実装されている。外側のマップがn-gramをキーに内側のマップを返し、内側のマップがそれに続くトークンとその出現頻度を保持する。

問題は、ドラフティングのたびに内側のマップが値渡しで不必要にコピーされていたことだ。TirmziはこれをC++の参照渡しに修正するPRを作成した。これだけで、コーパスサイズに応じてドラフティングが4.5倍〜25.6倍高速化された。

これほどの効果が出たのは、ドラフティングループが1トークン生成ごとに繰り返し走り、大きなマップのコピーが何度も発生していたためだ。コーパスが大きいほどマップも大きくなるため、541MBのフルコーパスで25倍という最大の効果が現れた。

残り3つの最適化

外側マップを`ankerl::unordered_dense`に置き換え

標準ライブラリのstd::unordered_mapはリンクリストによるチェイニングでコリジョンを解決するため、キャッシュフレンドリーでないことが広く知られている。TirmziはこれをMartin Ankerlのunordered_denseライブラリ、具体的にはsegmented_mapバリアントに置き換えた。

通常のmapバリアントは容量超過時に2倍に拡張するため、541MBフルコーパスではベースラインより1.16倍多くのメモリを消費した。segmented_mapは4096バイト単位でセグメントを追加する方式のため、この問題を回避できる。効果は静的キャッシュのロード時間を1.41〜1.65倍短縮、ドラフティングを1.02〜1.13倍高速化、メモリを1.07〜1.11倍削減だ。

内側マップをソート済みベクタ＋ブランチレス二分探索に置き換え

WikiText-103で構築した静的キャッシュを分析すると、64%の2-gramのフォロワー（続くトークン）は1種類のみだ。大半のn-gramに対してハッシュマップを維持するのはメモリの無駄であるため、ソート済みstd::vector<pair<token_id, count>>に置き換えてメモリ効率を改善した。

ただし一部の頻出2-gramは数千種類のトークンが後続するため、線形探索では遅くなる。そこでstd::lower_boundでO(log n)の二分探索を維持しつつ、さらにブランチレス実装に書き直した。標準の二分探索は次の探索範囲の計算が比較結果に依存するため、メモリ読み込みが完了するまでCPUパイプラインがストールする。Tirmziの実装はその依存を取り除き、CPUの投機実行を活かせる形にした。

// Tirmziによるブランチレス二分探索実装
const value_type * base = pairs;
while (n > 1) {
    const size_t half = n / 2;
    base = base[half].first < token ? base + half : base;
    n -= half;  // 比較結果に依存せず、CPUパイプラインのストールを回避
}
return (base - pairs) + (base->first < token);

静的キャッシュのシリアライズ形式の改善

静的キャッシュのファイルフォーマットも見直した。従来の形式はロード時に要素ごとに個別のメモリアロケーションが発生する構造だったが、新形式では連続したメモリ領域への一括シリアライズを採用した。これによりロード時のアロケーション回数が大幅に減り、ロード時間とメモリ使用量がさらに削減されている。

Daniel Lemireによる追加最適化で最終的に140倍へ

記事公開後、『Algorithms for Modern Hardware』の著者として知られ、SIMD命令や低レベル最適化の研究で著名なDaniel Lemireがさらなる最適化をPRとして送付した。具体的にはキャッシュのルックアップ処理へのSIMDベクタ化と、ハッシュ計算のさらなる効率化が含まれている。これによりTirmziの最適化に加えてさらに最大4.2倍の高速化が実現し、アップストリームのllama.cppと比較した総合的なスピードアップは最大140倍に達した。