推測によるSonarの加速

推測デコード（Speculative decoding）は、高速かつ小型のドラフトモデルを用いて完了候補を生成し、それをより大型のターゲットモデルで検証することで、大規模言語モデル（LLM）の生成速度を高速化する手法です。このスキームでは、コストの高いターゲットモデルが単一のトークンを生成する代わりに、単一のステップで複数のトークンを出力します。本稿では、PerplexityにおいてSonarモデルのトークン間レイテンシを削減するために適用した、さまざまな種類の推測デコードの実装詳細を紹介します。

推測デコード

推測デコードは、自然言語の構造とTransformerの自己回帰的な特性を利用して、トークン生成を高速化します。Llama-70Bのような大型モデルは、Llama-1Bのような小型モデルよりも多くの知識を保持していますが、単純なタスクにおいては同程度の性能を発揮します。この重複は、特定のシーケンスについてはコストの低いモデルで生成する方が効率的であり、複雑な問題は大型モデルに任せるべきであることを示唆しています。課題となるのは、どの補完がより優れているか、そして小型モデルによる生成が大型モデルと同等の品質であるかを判断することです。

幸いなことに、LLMは自己回帰的なTransformerであり、トークンのシーケンスが与えられると、次のトークンの確率分布を出力します。さらに、入力シーケンス内のトークンに関連付けられた中間特徴量から得られるロジットも、モデルがそれらのトークンを生成する確率を示しています。この特性により推測が可能となります。入力プレフィックスから開始して小型モデルがトークンシーケンスを生成した場合、それを大型モデルで実行してターゲットモデルとどの程度一致するかを判定できます。各候補のプレフィックスには確率によるスコアが付けられ、許容閾値を超えた最も長いものが選択されます。ボーナスとして、ターゲットモデルは後続のトークンも無料で提供するため、ドラフトモデルがn個のトークンを生成すれば、1回のステップで最大n + 1個のトークンを出力可能です。

拡大

推論時において、推測サンプリングプロセスはおおよそ以下の4つの段階に分けられます。

プリフィル（Prefill）：ターゲットモデルとドラフトモデルの両方を入力シーケンスで実行し、KVキャッシュエントリを埋める必要があります。Medusaのような手法では予測のために単純な密結合層を使用しますが、本稿では独自のKVキャッシュを必要とするTransformerベースのドラフトモデルに焦点を当てます。

ドラフト生成：ドラフトモデルが繰り返し処理を行い、一定数のトークンを生成します。ドラフトシーケンスは線形であるか、あるいはモデルが所定の深さまで木構造を探索することもあります（EAGLE、Medusa）。本稿では線形シーケンスに焦点を当てます。

受け入れ（Acceptance）：ターゲットモデルがドラフトシーケンスを実行し、各ドラフトトークンに対応するロジットを構築します。これにより、受け入れ可能な最長シーケンスの長さが決定されます。

ターゲット生成：ターゲットモデルがロジットを生成したため、不一致の位置またはシーケンスの末尾において、ロジットは後続のトークンに対応します。これらのロジットをサンプリングすることで、ターゲットモデルから堅牢なトークンを提供し、シーケンスを締めくくることができます。

推測デコードを実装するには様々な方法が存在します。本稿では、社内の1Bモデルを使用してSonarモデルを加速するために使用したスキームと、DeepSeek規模のモデルを高速化するために現在構築中の予測メカニズムに焦点を当てます。

ターゲット・ドラフト

推測デコードは、既存の小型LLMをドラフトモデルとしてターゲットモデルと組み合わせ、候補シーケンスを生成することで実現できます。本番環境では、ターゲットと同じデータセットでファインチューニングされたLlama-1Bモデルを使用してSonarを加速しました。このアプローチではドラフトモデルをゼロから学習する必要はありませんでしたが、小型モデルであっても相当量のKVキャッシュ容量を使用し、わずかなプリフィルオーバーヘッドを発生させるため、TTFT（Time To First Token）が増加します。

このスキームでは、デコーダーはデコード専用のバッチに対してのみ推測を行い、プリフィル時またはプリフィルとデコードが混在するバッチでは標準的なサンプリングを通じてトークンを生成します。プリフィル段階では、ターゲットのロジットが即座にサンプリングされ、ドラフトのKVキャッシュ内に新しく生成されたトークンもプリフィルされます。この時点ではドラフトはサンプリングされませんが、生成されたロジットはデコード段階に引き継がれます。

拡大

デコード段階では、ドラフトモデルが進行し、各段階でトップトークンをサンプリングします。目的のドラフト長に達した後、トークンがターゲットモデルを通過し、それに基づいてサンプラーが受け入れられたシーケンス長を特定するためのロジットを生成します。受け入れは、ドラフトとターゲットからの完全な確率分布を比較することで決定されます。ターゲットは常に受け入れられたドラフトシーケンスに続く1組のロジットを出力するため、それがサンプリングされて追加の出力が生成されます。ドラフトモデルはまだその受け入れられたトークンを認識していないため、次のデコードステップに備えて対応するKVキャッシュエントリを埋めるために再実行され、ロジットが再度引き継がれます。

EAGLE

EAGLEは、可能性の高いドラフトトークンの木構造の探索を通じて生成された、複数のドラフトシーケンスを調査する推測デコードスキームです。固定（EAGLE）または動的な形状（EAGLE-2）の木構造が、ドラフトトークンの連続的な実行を使用して探索されます。その際、線形シーケンスでスコアが最も高いトークンに従う代わりに、各ノードでTop-K候補が考慮されます。その後、シーケンスがスコアリングされ、最も適切なものが選択されて継続され、ターゲットからの追加トークンも付加されます。

拡大

より正確な予測を実現するために、EAGLEドラフトモデルはトークンだけでなく、ターゲットモデルの特徴量（最終層の隠れ状態）も使用して予測を行います。EAGLEの欠点は、低いレイテンシの予算内で適切な候補を生成できる、正確なカスタム小型ドラフトモデルを学習する必要があることです。一般的に、ドラフトモデルは元のモデルのデコーダー層と同一の単一Transformer層であり、埋め込みおよびlm_head投影と結びつけることでターゲットと密接に結合されています。これによりKVキャッシュ容量の消費が抑えられるため、EAGLEはメモリフットプリントが小さくなります。

ターゲットモデル内で木構造のシーケンスを検証するには、カスタムアテンションマスクを使用する必要があります。残念ながら、シーケンス全体にカスタムアテンションマスクを使用すると、現実的な入力長においてアテンションが大幅に（最大50%）低速化し、推測によって得られるはずの高速化の一部が相殺されてしまいます。このため、我々は完全な木構造探索を本番環境には未導入であり、DeepSeek-V3テクニカルレポートで提示されているMTPのようなスキームによる単一トークン予測の特殊なケースに注力しています。

MTP

このスキームは、予測にトークンと併せて隠れ状態を使用する点を除けば、ドラフト・ターゲットデコードと類似しています。通常のドラフト・ターゲット推測と比較して、プリフィル段階とデコード段階の両方でわずかに多くの作業が必要となります。ドラフトモデルはトークンと隠れ状態の両方を使用します。トークンt_{i+1}は、トークンt_iに対応するロジットL_iからサンプリングされ、これはさらに隠れ状態H_iから導出されます。その結果、入力トークンバッファは、ターゲットが出力する隠れ状態ベクトルに対して1ステップ左にシフトする必要があります。以下の図は、トレーニングで使用される対応関係と、推論中のシフトを示しています。

拡大

デコードフローはドラフト・ターゲットデコードと非常に似ていますが、隠れ状態とロジットの両方が引き継がれる点が異なります。我々の実装では関連するすべてのサンプリングとロジット処理を共有し、モデルの順伝播呼び出しのみを専門化しています。複数のトークンが予測される場合、ドラフトモデルは予測にドラフトの隠れ状態を使用し、自身の特徴量に基づいてKVキャッシュエントリを埋めます。長期的には、これが精度を低下させる可能性があります。その後、ターゲット予測のためにKVキャッシュエントリを埋めるべくドラフトモデルを実行する際は、より正確なターゲットの隠れ状態を入力としてシーケンス全体に対して実行します。これらのドラフトモデルは小さいため、追加のトークンを処理するコストは無視できる程度です。

MTPヘッドのトレーニング

MTPの利点を活用するために、我々はPerplexityのデータセット上でファインチューニングされたモデルに接続するMTPヘッドを、8xH100デバイスを搭載した1つのノードで学習するためのインフラストラクチャを構築しました。約1日で、Llama-1BからLlama-70B、DeepSeek V2-Liteに至るまでのモデルのヘッドを構築可能です。より大きなモデルについては、ファインチューニングプロセス中に構築されたMTPヘッドに依存しています。

MTPトレーニングの目標は、ターゲットの隠れ状態から外挿されたドラフトの隠れ状態とロジットを、ターゲットの次のトークンのロジットと隠れ状態に一致させることです。隠れ状態の推論はコストが高いため、トレーニング中に使用できるよう、推論に最適化されたターゲットモデルの実装を使用してそれらを事前に計算します。ただし、推論MTP実装を検証し、量子化や最適化による数値的な差異が結果を妨げないようにするため、検証損失と精度の推定にはターゲットモデルとドラフトモデルの両方の推論実装を完全に再利用します。

元の論文で使用されたShareGPTデータセットからより大規模なサンプルにスケールアップする際、EAGLE論文で概説および実装されたMTPヘッドアーキテクチャが、70Bサイズのモデルではトレーニングに失敗することに気づきました。短く多数のシーケンスを含むShareGPTとは異なり、我々はわずかに少ない数で、実質的に長いプロンプトを使用してトレーニングを行います。元のEAGLEヘッドは典型的なTransformerとは構造がわずかに異なっていたため、取り除かれていたRMS正規化層をいくつか再導入しました。その結果、トレーニングが収束するようになっただけでなく、ヘッドの精度も数パーセント向上しました。

拡大

層の正規化（Layer Norm）はトレーニングを促進するだけでなく、再導入は数学的にも直感的です。MTPヘッドはターゲットモデルの埋め込みとロジット投影を再利用します。これらは（Llama 70Bでは約2GBと）サイズが大きくなる可能性があるためです。トレーニング中、これらは凍結され、MTP層が元のモデルの投影層が学習したのと同じベクトル空間に予測を埋め込むことを学習すると想定されます。正規化を取り除くと、単一のMLPがMLPの後に正規化を続けた場合と同じ機能を学習することが期待されますが、これがドラフトモデルとターゲットモデルの隠れ状態の一致を妨げます。

推測デコードによる推論

推論エンジンにおいて、入力シーケンスのトークンを生成するには、まずそれらを適切なサイズのバッチにグループ化し、次に次のトークンのためにKVキャッシュ内にページを割り当てる必要があります。入力トークンとKVページ情報はバッファにパックされ、モデルを実行するすべての並列ランクにブロードキャストされます。最後に、メタデータがGPUメモリにコピーされ、モデルが実行されて次のトークンをサンプリングするためのロジットが生成されます。

リクエストを調整するラッパーを介してドラフトとターゲットの推論サーバーを緩やかに結合する特定の実装とは異なり、我々のドラフト・ターゲットペアは密接に結合されており、生成ステップを同期して進行します。バッチスケジューリングとKVページの割り当ては、すべての形式の推測デコード間でモデル間で共有されます。これにより、すべてのモデルが同じインターフェースを公開するため、モデルと包括的な推論サーバーを橋渡しするロジックが統一されます。

Perplexityの推論ランタイムはFlashInferを中心に形成されています。これはアテンションカーネルを設定およびスケジュールするために構築する必要があるメタデータを決定します。プリフィル、デコード、または検証のためにバッチを形成する入力シーケンスが与えられると、中間バッファを割り当て、アテンションで使用される特定の定数バッファを埋めるためにCPU側の作業を行う必要があります。この作業は、バッチスケジューリングとKVページ割り当てのコストに加えて行われるものであり、これらもGPUの稼働率を最大化するために隠蔽すべきレイテンシを生じさせます。

推測なしの推論についてはCPU側とGPU側の作業を完全に並列化しましたが、推測デコードのためのCPUとGPUのバランスはより複雑であることがわかりました。主な課題は、受け入れられたトークンの数が後続の実行のシーケンス長を決定し、回避困難なGPUとCPUの同期ポイントが生じるという点です。我々はCPU作業のレイテンシを最適に隠すために、さまざまなスケジューリングスキームを実験しました。

ドラフト・ターゲットのスケジュール

ターゲットモデルよりも小さいとはいえ、LLM全体をドラフトとして使用すると、GPU上でかなりのレイテンシが発生し、高コストなCPU操作を隠すための余裕が生まれます。小型モデルはテンソル並列処理の恩恵を受けられないため、ターゲットとドラフトがシャーディングされるランク数に不一致が生じます。我々の実装では、ドラフトモデルはTPグループのリーダーランクでのみ実行されます。

拡大

前述の通り、デコードステップはロジットを次の実行に引き継ぎます。これにより、ドラフトモデルの1回の実行をCPU側のバッチスケジューリング作業とオーバーラップさせることができます。バッチが組み立てられた後、サンプラーとドラフトを繰り返し呼び出すことでドラフトトークンが生成されます。並行して、検証用のバッチがターゲットモデルのために組み立てられ、並列ワーカーと同期されます。ターゲットのロジットが検証・サンプリングされ、受け入れられたシーケンス長が決定されます。この時点で、後続のシーケンス長を決定するためにGPUとCPUの同期が必要となります。ドラフトモデルはリーダーノードでのみ実行されるため、そのバッチは逐次的に設定され、ターゲットが生成した追加トークンでKVキャッシュエントリを埋めるために実行が開始されます。現在のドラフト実行によって生成されたロジットは、後続の実行における最初のドラフトトークンをサンプリングするために使用されます。最も重要な点は、ドラフトが実行されている間に、次のバッチをスケジュールできることです。

単一トークンのためのMTPスケジュール

ランタイムはまだEagle形式のドラフト木探索を提供していませんが、単一のTransformerデコーダー層のサイズのモデルによって生成されるドラフトトークンの線形シーケンスを考慮する、このスキームの特殊なケースを実装しました。このスキームは、DeepSeek R1のオープンソース重みを使用したドラフト予測に使用できます。単一トークンを予測するサブケースは、大規模なMTP層がオーバーヘッドを正当化するのに十分な高い受け入れ率を達成するため、興味深いものです。

MTPスケジューリングはやや複雑です。ドラフトモデルがはるかに高速であり、CPU側のレイテンシを隠す余地が少ないためです。さらに、ドラフトはターゲットモデルと一緒にシャーディングされるため、バッチ情報のために共有メモリ転送が必要になります。実行は、前述のスキームと同様に、バッチ情報を転送し、引き継いだロジットから最初のトークンをサンプリングすることから始まります。次に、ターゲットが実行されてトークンを検証し、2 * Dトークンを処理します（Dはデコードバッチサイズ）。バッチが2つの半分に均等に分割されるため、InfiniBandのような低速なインターコネクトを介したMixture-of-Experts（MoE）モデルのマイクロバッチ処理には理想的です。ターゲットの隠れ状態は次のドラフト実行に引き継がれ、ロジットは検証のためにサンプラーに渡されます。

拡大

GPU上で限られた追加作業を行うことで、ドラフトシーケンスの受け入れ後のCPUとGPUの同期を回避します。ターゲットの入力トークンがシフトされた後、カーネルが次のターゲットトークンを対応する位置に挿入します。その後、ターゲットと同じバッチ情報でドラフトが再実行され、KVキャッシュエントリを埋め、次回の実行のためのロジットと隠れ状態を構築し、受け入れられなかったトークンに対しては冗長な作業を行います。このような状況では、ドラフトモデルのサイズが小さいため、未使用の作業によるレイテンシはほとんど測定できません。ドラフト実行と並行して、CPU上でシーケンス長が決定され、GPU作業の終了を待つことなく次のバッチのスケジューリングが開始されます。

ドラフト層における追加作業のオーバーヘッドはアテンションでは目立ちませんが、MLP層はより問題となります。行列乗算命令はトークン数の次元に沿って64の境界までパディングされるため、倍増させても必要なブロック数が大幅に増えない場合は、オーバーヘッドは隠蔽されます。より長いドラフトシーケンスの場合、オーバーヘッドのコストが高くなるため、通常のドラフト・ターゲットモデルで使用されるスキームの方がうまく機能します。

参考文献

推測デコードによるTransformerの高速推論

EAGLE: 推測サンプリングには特徴量の不確実性の再考が必要

EAGLE-2: 動的ドラフト木による言語モデルの高速推論

EAGLE-3: 学習時テストによる大規模言語モデルの推論加速のスケールアップ

Medusa: 複数のデコードヘッドを備えたシンプルなLLM推論加速フレームワーク

層正規化（Layer Normalization）

FlashInfer: LLM推論サービングのための効率的かつカスタマイズ可能なアテンションエンジン

DeepSeek-V3テクニカルレポート

引用

@online{perplexity2025AcceleratingSonarThroughSpeculation,
  author = {{Perplexity} and Nandor Licker and Erin McKnight},
  title = {{Accelerating Sonar through speculation}},
  date = {2025-06-10},
  url = {https://perplexity.coolkeypoint.com/hub/blog/accelerating-sonar-through-speculation}
}