AMD GPU上のvLLMにおける投機的デコーディング

目次TL;DR: 投機的デコーディングにより、vLLMは単一のターゲットモデルパスで複数のドラフトトークンを検証できます。私たちの実験では、出力トークンスループットへの影響は、ドラフト手法や提案長によって異なり、モデルファミリー、ドラフトチェックポイント、ワークロード、および受け入れ動作にも依存しました。はじめに大規模言語モデルは幅広いアプリケーションをサポートしますが、それらを大規模に提供するには慎重な最適化が必要です。標準的な自己回帰デコーディングは、ほとんどのLLMサービングシステムで使用されるベースラインです。モデルは1つのトークンを生成し、それをシーケンスに追加し、更新されたシーケンスを使用して次のトークンを生成します。このプロセスはシンプルで信頼性がありますが、出力トークンは厳密な左から右の順序で生成する必要があるため、サービングループは依然として1つのコミット済みトークンずつ進みます。投機的デコーディング[1]は、ドラフト・アンド・検証メカニズムを通じてこのベースラインを拡張します。軽量なドラフトコンポーネントが候補トークンを提案し、ターゲットモデルがそれらがコミットされる前にそれらの候補を検証します。複数のドラフトトークンが受け入れられると、システムはターゲットモデルの出力動作を維持しながら、単一のターゲットモデル検証ステップで複数の出力トークンをコミットできます。この記事では、vLLMで投機的デコーディングがどのように機能するかを調査し、テスト環境からの測定結果を共有します。まず、自己回帰デコーディングのベースラインとドラフト・アンド・検証プロセスをレビューします。次に、ネイティブMTP、Gemma 4 MTP、EAGLE-3、DFlash、DSparkの5つの投機的ドラフトアプローチを調べます。これらの手法は、ドラフトコンポーネントがターゲットモデルからどのように情報を受け取るか、および候補トークンが逐次的に、自己回帰的に、並列に、またはハイブリッドアプローチで生成されるかによって異なります。最後に、テスト環境でこれらの手法を有効にする方法を示し、ROCm™オープンソフトウェアプラットフォームを使用したAMD Instinct™ MI300XおよびMI355X GPUでの実験からの測定結果を報告し、実用的なチューニングとオブザーバビリティの考慮事項について議論します。自己回帰デコーディングのベースライン標準的な自己回帰デコーディングでは、各デコードステップは1つの新しいトークンを生成してコミットします。例えば、4つの出力トークンを生成するには、4つの逐次的なデコードステップが必要です。ステップ1：コンテキスト → モデル → T1ステップ2：コンテキスト + T1 → モデル → T2ステップ3：コンテキスト + T1 T2 → モデル → T3ステップ4：コンテキスト + T1 T2 T3 → モデル → T4各ステップの後、生成されたトークンはシーケンスに追加され、次のステップの入力の一部となります。これによりデコーディングループは単純になりますが、出力トークンごとに1回のモデルデコードステップが必要になります。長い生成の間、このトークンごとのループはレイテンシを支配し、サービングスループットを制限する可能性があります。したがって、投機的デコーディングの鍵となる疑問は次のとおりです。生成が一度に1つのトークンずつしか進まない頻度を減らしながら、元のモデルの出力動作を維持できるか？投機的デコーディングは、提案と検証を分離することによってこれを解決します。ドラフトコンポーネントが最初に複数の候補トークンを提案します。元のモデルが、それらがコミットされる前にそれらの候補を検証します。投機的デコーディングのコアアイデア投機的デコーディングは元のモデルを置き換えるものではありません。代わりに、元のモデルをターゲットモデルとして維持し、最終出力の責任を負わせ、その前に高速な提案ステージを追加します。プロセスは2つの部分からなります。ドラフト：複数の候補トークンを提案します。検証：ターゲットモデルを使用してそれらの候補をチェックします。各投機的デコーディングラウンド中、図1に示すように、軽量なドラフトコンポーネントが1つ以上の将来のトークンを提案します。これらのトークンは候補にすぎず、すぐにコミットされるわけではありません。次に、ターゲットモデルが候補トークンシーケンスを1回の検証パスで評価します。検証は左から右に進みます。各ドラフトトークンは、対応する位置でのターゲットモデルの結果を使用してチェックされます。受け入れられたトークンは出力シーケンスにコミットされます。ドラフトトークンが拒否された場合、同じ提案からの後続の候補はもはや受け入れられません。ドラフトトークンが拒否された場合、ターゲットモデルが次のトークンを提供します。残りのドラフトトークンは破棄され、更新されたシーケンスから生成が続行されます。概念的には、標準的な自己回帰デコーディングは次のように進みます。ターゲットモデル →T1ターゲットモデル →T2ターゲットモデル →T3ターゲットモデル →T4投機的デコーディングは、代わりに複数の候補位置を一緒に評価することを可能にします。ドラフトがT1 T2 T3 T4を提案します。モデルが✓ ✓ ✗停止を検証します。T1 T2をコミットします。置換トークン-これにより、複数の候補が受け入れられた場合にターゲットモデルのデコードラウンドの数を減らすことができます。ドラフトコンポーネントがターゲットモデルに受け入れられるトークンを生成すると、単一のターゲットモデル検証ステップから複数の出力トークンをコミットできます。提案が拒否された場合、ターゲット側の結果が生成の続行方法を決定します。単純な受け入れ/拒否の例図2は、1回の投機的デコーディングラウンドの例を示しています。緑色のボックスは検証を通過したドラフトトークン、赤いボックスは最初に拒否されたドラフトトークンを示し、灰色のボックスは破棄された後続のドラフトトークンです。出力の青いトークンは、ドラフト提案からではなく、ターゲットモデルからのものです。現在のプロンプトが次のとおりだと仮定します。今日の天気はドラフトコンポーネントがいくつかの将来のトークンを提案します。晴れて暖かいです。ターゲットモデルは左から右にドラフトトークンを検証します。ドラフトが晴れて暖かいです。と提案します。モデルが✓ ✓ ✗停止を検証します。晴れて、とコミットします。最初の2つのドラフトトークン、「晴れて」と「、」は受け入れられます。3番目の位置では、ドラフトは「暖かい」を提案しますが、ターゲットモデルは「です」を選択します。後続の候補「外」は、最初の拒否された位置に続くため破棄されます。したがって、次のデコーディングラウンドは次から続行されます。今日の天気は晴れて暖かいです。ドラフト手法の仕組みすべての投機的デコーディング手法は同じ全体的なドラフト・アンド・検証プロセスに従いますが、ドラフトコンポーネントの設計方法やターゲットモデルとの連携方法が異なります。主な違いは次のとおりです。ターゲットモデルから受け取る情報の種類。この情報がドラフトプロセスにどのように組み込まれるか。候補トークンが逐次的に生成されるか並列に生成されるか。これらの違いに基づいて、この記事で議論されているドラフト手法は、3つの広範なカテゴリに分類できます。ネイティブMTPモジュール、独立したMTPドラフター、および専用のターゲット条件付きドラフトネットワーク。ネイティブMTPモジュール：ターゲットモデルアーキテクチャに直接組み込まれています。モデルネイティブの補助予測パスを使用します。候補トークンを逐次的に生成します。独立したMTPドラフター：特定のターゲットモデルとペアになった独立したチェックポイントを使用します。推論中にターゲットモデルの活性化と共有KVキャッシュ情報を使用します。候補トークンを逐次的に生成します。専用のターゲット条件付きドラフトネットワーク：EAGLE-3、DFlash、DSparkなど、特定のターゲットモデル用にトレーニングされた独立したスペキュレーターモデルを使用します。EAGLE-3はターゲットモデルの隠れ状態から自己回帰的にドラフトし、DFlashはターゲットモデルの隠れ状態から並列ブロックをドラフトし、DSparkは軽量な因果補正と信頼度ベースのプレフィックス選択を追加します。これらのカテゴリはドラフトコンポーネントのアーキテクチャを表しており、ターゲットモデルファミリーを表すものではありません。ターゲットモデルは、ネイティブMTPをサポートすると同時に、独立してトレーニングされたEAGLE-3、DFlash、またはDSparkドラフトモデルを持つこともできます。ドラフトコンポーネントは完全に単独で動作するわけではありません。手法によっては、ドラフトコンポーネントは次の情報を受け取ることがあります。ターゲットモデルからの隠れ表現。選択された複数のターゲットレイヤーからの隠れ状態。ターゲットモデルのKVキャッシュ。複数のターゲットモデル表現を組み合わせることで生成される特徴。次のセクションでは、各手法がこの情報をどのように使用し、候補トークンをどのように生成するかを説明します。ネイティブMTPマルチトークン予測（MTP）は、即時次のトークンを超えるトークンを予測するためのモデルネイティブメカニズムのファミリーを指します。vLLMでは、ネイティブMTPは利用可能です。