SageMaker HyperPod でのLLM推論のための分離型プリフィルとデコード

大規模言語モデルの推論には、プリフィルとデコードという2つの根本的に異なるフェーズがあります。プリフィルはコンピュートバウンドであり、入力プロンプト全体を並列に処理して初期のKVキャッシュを生成します。デコードはメモリバウンドであり、一度に1トークンを生成し、モデル重みと増大するKVキャッシュにアクセスするために大量のメモリ帯域幅を必要とします。プリフィルとデコードが同じGPUを共有すると、長いプロンプトがすべての同時リクエストのトークン生成を停止させます。分離型プリフィルとデコード（DPD）は、各フェーズをElastic Fabric Adapter（EFA）で接続された別々のGPUプールで実行することにより、この干渉を除去します。これにより、ファーストトークン時間（TTFT）とトークン間レイテンシ（ITL）を独立に調整でき、テールレイテンシをより確実に制御し、長コンテキストのプリフィルが進行中のデコードリクエストをブロックするのを防ぎます。vLLMは連続バッチ処理とPagedAttentionによって単一ノードの効率を向上させますが、大規模に展開する組織はマルチノードのオーケストレーションとルーティングの最適化において依然として課題に直面しています。

この記事では、HyperPod推論オペレーターを使用してAmazon SageMaker HyperPod上でvLLMを用いたDPDを実装する方法を示します。

分離型推論を使用するタイミング

プリフィルとデコードの分離は、長コンテキストで高同時実行のストリーミングワークロード（チャットアシスタント、エージェントパイプライン、ドキュメント分析エンドポイント、大規模な検索コンテキストを用いたRAGなど）に最大の効果を発揮します。これらのケースでは、同じGPU上の長いプロンプトが他のすべてのリクエストのデコードを停止させ、トークンあたりのレイテンシスパイクを引き起こしますが、DPDは構造的にこれを除去します。入力プロンプトが定期的に4096トークンを超える場合、複数の同時ユーザーまたはリクエストがある場合、ストリーミング応答で一貫したトークン配信が重要な場合、長いプロンプトと短いプロンプトが混在するトラフィックがある場合にDPDを検討してください。GPU競合が問題にならないバッチ処理やオフラインワークロード、低同時実行デプロイメント、短いプロンプトのみのトラフィックでは、同じ場所でのデプロイメントがより簡単な選択です。ルーティングしきい値を下回る場合、EFA RDMAを介したKVキャッシュ転送の固定コストがデコード分離の利益を上回ります。DPDルーターはこれらのリクエストを直接デコーダーに送信するため、単一のエンドポイントで長短混在トラフィックを自動的に処理できます。

DPDには、RDMA対応のEFAネットワークを備えた少なくとも1つのプリフィルノードと1つのデコードノードが必要です。サポートされるインスタンスタイプについては、DPDモデルエンドポイントのデプロイセクションを参照してください。

アーキテクチャ

HyperPod DPD実装はvLLMプロダクションスタックルーター上に構築され、LMCacheがNIXLとEFAを介したKVキャッシュ転送層を提供します。デプロイメントには3つのコンポーネントとトランスポートスタックがあります。

• インテリジェントルーター: コントロールプレーン。各プロンプトをトークン化し、設定可能なトークンしきい値を適用して、リクエストが分離パスを取るかエンドツーエンドでデコーダーに送られるかを決定します。
• プリフィラーポッド: LMCacheをKVコネクタとして使用するvLLMワーカー。長いプロンプトのKVキャッシュを計算し、選択されたデコーダーにプッシュします。
• デコーダーポッド: 受信側のvLLMワーカー。GPUメモリを予約してKV転送を受け入れ、転送完了と同時に生成を開始します。
• KV転送: LMCache PD、NIXL、libfabric、EFAの4層スタックを使用し、転送コストはプリフィル計算に比べて無視できる程度です。

デプロイメントの概要

DPDはHyperPod推論オペレーターによって実装される機能です。このセクションでは、前提条件、推論オペレーターのインストール、およびLlama 70Bモデルを効率的に提供するDPDを使用した推論エンドポイントのデプロイメントについて説明します。前提条件として、AWS CLI、HyperPodクラスターへのkubectlアクセス、HuggingFaceトークン、十分なサービス割り当てが必要です。HyperPod推論オペレーターバージョン3.2以降が必要であり、新しいHyperPod EKSクラスターにはデフォルトでインストールされています。

この例では、2つのml.p5.48xlargeインスタンスにMeta Llama 3.3 70Bモデルをデプロイします。DPD推論デプロイメントでは、NVLinkとEFAの両方をサポートするインスタンスタイプを選択します。ワーカーイメージにはvLLM、LMCache、NVIDIA NIXL、およびEFA libfabricプロバイダーが含まれている必要があります。サポートされるイメージオプションは、オープンソースLMCacheまたはSageMaker Deep Learning Containerです。

デプロイメントマニフェストの主要なDPD関連フィールドには、spec.pdSpec（プリフィル/デコードトポロジーの宣言）、routingThreshold（ルーティングしきい値）、および各ロールのargsがあります。マニフェストを適用すると、オペレーターがプリフィル用とデコード用の2つのDeploymentオブジェクトと、ルーター用のDeploymentを作成します。