分散されたプリフィルおよびデコード

入力プロンプトから出力トークンを生成するには、LLM推論がプレフィルとデコードの2段階に分けられます。プレフィルは入力トークンで動作し、KVキャッシュを充填した後、1つずつトークンを生成するデコード段階に入ります。

単一のデコードステップは通常数十ミリ秒で実行されますが、プレフィルは大幅に時間がかかります。同じデバイスで実行されると、プレフィルとデコードを混ぜるとデコードのパフォーマンスが低下します。この記事では、プレフィルとデコードを分離して異なるデバイスで実行し、プレフィルスループットとデコードの遅延を最大化するソリューションについて探ります。

プレフィル vs デコードパフォーマンス

一般的なLLMサービングエンジンでは、バッチスケジューラがモデルの各実行ステップで処理するリクエストを選択します。単一のデバイスまたはノードで実行される場合、プレフィルとデコードの両リクエストは一緒にバッチ処理されます。系列長に沿って集約されるアテンションのコストは、KVキャッシュのエントリー長（kv_len）に比例してプレフィルとデコードの両方で増加します。デコードリクエストは通常、系列のトークンを独立して処理する他の層を通過する際、単一のトークンを前進（qo_len=1）するだけです。プレフィルリクエストは数千または数万のトークンを高コストで密な層を通過します（qo_lenが大きい）。

フォワードパスのレイテンシは、主に密な層を通過する独立トークン数（qo_len）によって影響され、 KVキャッシュから取得されたトークン数（kv_len）より影響を受けにくいです。アテンションはリクエスト数と系列長に対するkv_lenと並行化でき、良好な利用化を実現します。プレフィルは計算量に依存しています： qo_len が高く、GEMMカーネルはモダンGPUの計算能力をフルに活用するためにM軸に沿って十分なブロックを割り当てることができます。デコードはメモリに依存しています：通常の低いバッチサイズにより、M軸に沿った入力数が通常小さく、一ブロックにしか効かない。SR-K GEMMカーネルは低トークンバッチサイズのSM利用率を改善できるが、キャッシュと行列積ユニットは通常、十分に活用されていません。

混在する場合、プレフィルリクエストを含むバッチはフォワードパスを通過する際により高いレイテンシを引き起こし、インスタンス全体のデコードスループットに悪影響を与えます。プレフィルリクエストをデコードリクエストと混合させたり、チャンク化したプレフィルを採用してデコード性能をわずかに改善することはできますが、プレフィルスループットを十分に維持して、インスタンス上の十分なリクエストを処理しない限り、デコードスループットを最大化することは難しいです。大規模モデルの場合、典型的な出力の長さで、デコードのために大きなバッチサイズを保つためには、プレフィルが頻繁に行われ、結果として平均レイテンシが著しく低下し、出力にスタッタを引き起こします。

これらの問題は、プレフィルとデコードを実行するための別々のノードセットを使用することで対処できます。プレフィラーノードを複数のデコーダーノードと関連付けることで、プレフィルするための十分なリクエストをスケジュールし、デコーダーノードでの同時リクエスト数も最大化することが可能です。プレフィラーノードはKVキャッシュを充填し、それらをデコーダーノードに転送します。デコーダーはもうプレフィルのために停止する必要がないため、レイテンシははるかに決定論的になり、アクティブなリクエストのkv_lenの成長の全体的な影響もはるかに少なくなります。そのコストは、KVキャッシュのネットワーク転送が数十から数百ミリ秒かかる可能性があるため、First Tokenへの時間（TTFT）の増加に支払われます。

KVメッセンジャー

Perplexityでは、プレフィルとデコードの分離型実装は、プレフィラーノードからデコーダーノードへのKVキャッシュ転送をネットワークを通じて調整するKVメッセンジャーを中心に構築されています。プレフィラ側では、メッセンジャーがデコーダーノードからのリクエストを受け入れ、それをバッチスケジューラに渡し、フォワードパスの実行を追跡してできるだけ低いレイテンシでKVキャッシュを送信します。デコード側では、未退避ページが割り当てられた後、メッセンジャーはKVキャッシュとデコーダーコンテキストの転送が完了することを通知されるまでリクエストをデコードのためにスケジュールすることをブロックします。

プレフィルを分離するには高スループット、低レイテンシ接続が必要なため、我々の実装はRDMAに合わせて設計されており、EFAとConnectXネットワークインターフェースコントローラ（NIC）の両方をサポートしています。libfabricに基づいたKVメッセンジャーは、fabric-libラッパーを使用してリモートダイレクトメモリアクセス（RDMA）プリミティブに対する高次低レイテンシ抽象化を提供し、効率的なページとメタデータ転送、および低レイテンシシグナルを実装しています。バックグラウンドでfabric-libはプレフィラーノードからデコーダーノードへのデータのコピーを行うためにGPUと直接接続されたNICを調整します。

受信後、プレフィラーノードは対応するソースKVページセットを割り当て、ローカルエンジンを使用してプレフィルのリクエストをスケジュールします。レイテンシを最小限に抑えるため、転送はフォワードパスを待たず、モデルが個々の層のKVキャッシュエントリーをKVキャッシュに追加し終えるとすぐにKVページのコピーが開始されます。プレフィルリクエストはチャンク化できるため、バッチスケジューラは実行前に現在スケジュールされているチャンクをKVメッセンジャーに通知します。CUDAグラフをサポートしながらレイヤを追跡できるようにするため、メッセンジャーは各レイヤの出力プロジェクションの後にインクリメントされるカウンターをポーリングする専用スレッドを保持しています。カウンターはシャーディング環境でリードノードのみに維持されます： KVキャッシュエントリは追加後かつアテンション前に有効であるが、出力プロジェクションはランク間で縮小され、それらを暗黙に同期します。カウンターの変化が観察されると、メッセンジャーに通知され、fabric-lib を呼び出してレイヤの転送を開始します。

最後のチャンクの転送が完了した後、追加のメタデータもコピーされます： 推測デコーディングまたはMTPはロジットと隠れ状態をデコーダーに移動させる必要があります。これらのコピーもRDMAを通じて事前に割り当てられたバッファーで実行されます。

最後のチャンクの保留中の全転送の完了後、プレフィラーノードはKVページを解放し、リクエストを完了させます。デコーダーノードは明示的に通知されません： 代わりに、完了した操作の数を追跡するために即時カウンターを使用します。プレフィラ側のRDMA操作の数は転送されたページ数に比例します。既知のページとコンテキストコピーの数が完了すると、fabric-libはKVメッセンジャーにデコードの準備ができたことを示唆します。メッセンジャーはコンテキストを解放し、リクエストをLLMエンジンに引き渡します。

シャード化KVキャッシュ転送

プレフィラとデコーダーがテンソル並列処理（TP）に依存し、KVキャッシュを同一にシャーディングまたは複製する場合、単一の転送エンジンが複数のデバイスを調整して全レプリカのページを送信および受信します。モデルの実行プログラムが複数のデバイスやプロセスにレプリケートされていても、単一のメッセンジャーと転送エンジンを使用するために、cuMemおよびcuMemImportFromShareableHandleを使用して、KVキャッシュをサポートするデバイスメモリーを割り当て、メインプロセスにマッピングします。転送エンジンはノードのトポロジーを検査し、各KVキャッシュスライスの転送に使用する最寄りのNUMAノードにあるNICとCPUを見つけます。

送信元と宛先が同一にシャーディングされている場合、転送は簡単です。送信元と宛先のデバイスとページに一対一のマッピングがあるためです。この状況では、シャーディングは転送レイテンシを暗黙に助けます： より多くのGPUを使用することで、より多くの関連NICが活用され、帯域幅の完全利用に近づきます。しかし、不一致がある場合、転送エンジンは送信元と宛先のスライス間の比率に応じてページを分割または再構築する必要があります。

プレフィラーがより多くのデバイスにKVキャッシュを分割する場合、デコーダーでページを再構築し、対応する半分をプレフィラーデバイスから送信します。デコーダーがより多くのシャードを持つ場合、複数の送信元からページを受信します。デコーダーはプレフィラーのシャーディングスキームを知る必要があり、期待されるRDMA書き込み数を計算することができます。レプリケーションが関係する場合、プレフィラーはデバイスをレプリカセットにグループ化し、プレフィラー自身で完全なKVキャッシュを複製します。宛先レプリカセットは送信元セットの1つをランダムに割り当て、すべての利用可能なデバイスを使用してRDMA書き込みを開始します。

シャーディングされた転送はKVキャッシュに微調整を必要とします。デフォルトでは、FlashInferはNHDレイアウトを依存し、ページ内のトークンをヘッド内に配置します。キャッシュは最も注意ヘッド数に沿ってシャーディングされている可能性が高いため、これはヘッド内で不連続性を生み出します。RDMA転送はストライド書き込みを暗黙にサポートせず、転送を実行する各ヘッドに対して1つの操作が必要です。代わりに、libfabricとのやり取りの数を減らすため、我々はKVキャッシュをHNDレイアウトを使用して配置し、ヘッドのディメンションをトークン数の前に置きます。これによりページ全体を1回の書き込みでコピーできる継続性が保証されます。

推測デコード

推測デコードは、分離されたプレフィルデコードにわずかな調整が必要です。我々の実装では、プレフィラーノードはトークンをサンプリングすることを許可されていません。PerplexityのSonarモデルは構造化出力をサポートしているため、プレフィラとデコーダー間でスキーマプロセッサ実装を同期させる複雑さを避けたいです。MTPおよび推測デコードメカニズムにおいて、ターゲットモデルからトークンをサンプリングすることがプレフィルのドラフトモデルまでの最終トークンを含まないようにします。

これらの問題を解決するために、プレフィルは入力シーケンスの最終トークンを含みません。代わりに、プレフィルで最終トークンの前に続く隠れ状態またはロジットを転送し、次のステップでデコーダーのデコードトークンとして扱います。これにより、プレフィル後に完全なデコードステップを実行して最初のトークンを発行する必要があるため、遅延がやや増加しますが、実装の複雑さは大幅に削減されます。

分離デプロイメント

我々は、プロダクションのトラフィックフォローまたは社内評価ワークロードをサポートするために、異なるモデルで複数の分離型構成を導入または実験しました。モデルのサイズと注意メカニズムに基づいて、プレフィラおよびデコーダーノードに適したシャーディングスキームを選択し、GPUを最大限に利用しました。

DeepSeek-R1

DeepSeekでは、テンソル並列（TP）およびデータ並列（DP）デプロイメントの両方を検討しました。以前のブログ投稿で議論したように、TPデプロイメントは低いスループットのコストで優れたレイテンシを提供し、重いトラフィックを提供するためにより多くのGPUを必要とします。DPデプロイメントは負荷とともによりよくスケールしますが、デバイス間またはノード間通信のコストのためにスループットのピークが低いです。

DeepSeekはマルチヘッド潜在注意機構に依存しており、KVキャッシュを圧縮します。すべてのKVヘッドが単一の潜在ベクターに圧縮されるため、TPはKVキャッシュをシャーディングできません。代わりに、すべてのランクで潜在ベクターを複製します。シャーディングは展開後に行われ、各ランクが同じ潜在表現から異なるヘッドを抽出できます。したがって、すべてのKVキャッシュシャードはプレフィラとデコーダシャードの間で同一です。

イントラノードTPセットアップでは、プレフィラとデコーダの両方が同じようにシャーディングされています。利用可能なすべてのNICを完全に活用するために、すべてのランクから転送が行われます。しかし、データ並列デプロイメントの場合、TPランクサイズが小さいか、各DPランクが単一GPUに割り当てられている場合、KVキャッシュの複製コピーを持つプレフィラデバイスはそれを転送できます。利用可能なすべてのNIC間でリクエストをバランスさせるために、プレフィラからデコーダへのKVキャッシュを送信するためにランダムにGPUとNICを選択します。

プレフィルとデコードを組み合わせた際、我々のR1デプロイメントは多数のミリ秒単位のプレフィルによる頻繁な中断のため、一貫して50 TPSを超えるのに苦労していました。一方、プレフィルを分離することで、各リクエストに対して約100msのTTFTペナルティが生じましたが、単一のプレフィラーノードが3つのデコーダーノードで一貫したバッチサイズを維持することができ、90 TPSを超えるスループットを提供しながらデコーダーノードあたり約1 QPSの負荷を処理できました。データ並列デプロイメントでは、TPSは50前後と若干低かったですが、インスタンスはランクあたり1 QPSの負荷を処理でき、1ノードには8つのランクがありました。

Qwen3-Coder

この480Bモデルは、Grouped-Query Attention（GQA）を使用しているため、注意は簡単にシャーディングされ、KVキャッシュを犠牲にすることなくテンソル並列処理の利点を享受できます。その結果、プレフィルとデコードの両方で8つのGPUにモデルをシャーディングし、1つのプレフィラーノードに約3つのデコーダーノードをペアリングすることができました。注意がシャーディングされているため、我々はプレフィラとデコーダのKVキャッシュを分割するためにHND KVキャッシュレイアウトに依存し、プレフィラランクをデコーダランクとペアリングし、すべてのNICを完全に活用してスライスを並行して転送します。