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

GPU でプレフィルとデコードを共有すると、長いプロンプトにより、すべての並行リクエストに対してトークン生成が停止してしまいます。Disaggregated Prefill and Decode (DPD) は、Elastic Fabric Adapter (EFA) を介して接続された別々の GPU プールで各フェーズを実行し、Remote Direct Memory Access (RDMA) を利用することで、この干渉を解消します。

大規模言語モデル（LLM）の推論には、根本的に異なる 2 つのフェーズがあります。プレフィルは計算集約型です。これは、入力プロンプト全体を並列処理して、初期のキーバリュー (KV) キャッシュを生成します。デコードはメモリー集約型です。これは一度に 1 トークンを生成し、モデル重みと成長する KV キャッシュへのアクセスには大量のメモリー帯域幅が必要です。

これらを専門的なエンジンに分離することで、各フェーズに異なる並列戦略を割り当てることができます。この分離により、最初のトークンまでの時間 (TTFT) とトークン間のレイテンシ (ITL) を独立して調整でき、チャンクドプレフィルのチューニングよりも信頼性が高くテールレイテンシを制御し、長いコンテキストを持つプレフィルが進行中のデコードリクエストをブロックすることを防ぐことができます。

vLLM は、連続バッチ処理と PagedAttention によって単一ノードの効率を向上させます。しかし、大規模展開を行う組織は、マルチノード展開のオーケストレーションやルーティングの最適化において依然として課題に直面しています。

本稿では、HyperPod Inference Operator を使用して Amazon SageMaker HyperPod 上で vLLM とともに DPD を実装する方法を示します。

分離型推論の適用タイミング

プレフィルとデコードを分離することで、長文コンテキストや高同時実行ストリーミングワークロードにおいて最大の効果が得られます。具体的には、チャットアシスタント、エージェントパイプライン、ドキュメント分析エンドポイント、および大規模な検索コンテキストを持つ検索拡張生成（RAG）などが該当します。これらのケースでは、単一の長いプロンプトが配置された GPU で処理される際、他のすべてのリクエストに対する実行中のデコードが停止し、トークンごとのレイテンシの急上昇を引き起こします。DPD はこの問題を構造的に解消します。

DPD の適用を検討すべきワークロードとは:

• 4,096 トークンを定期的に超える入力プロンプトを持つ場合。

• 複数の同時実行ユーザーまたはリクエストが存在する場合。

• トークンの一貫した配信が重要なストリーミングレスポンスを必要とする場合。

• 長いプロンプトと短いプロンプトの両方を含む混合トラフィックの場合。

GPU の競合が問題にならない場合は、配置型（colocated）デプロイの方がシンプルです。これは、TTFT（Time To First Token）を最適化するバッチ処理やオフラインワークロード、低同時実行のデプロイメント、または短いプロンプトのみを対象としたトラフィックに適しています。ルーティング閾値以下では、EFA RDMA を介して KV キャッシュを転送する固定コストが、デコードを分離することによるメリットを上回ります。DPD ルーターはこれらのリクエストを直接デコーダーへ送信します。その結果、単一のエンドポイントで長いプロンプトと短いプロンプトの混合トラフィックを自動的に処理でき、手動でのルーティングロジックは不要となります。

DPD には、少なくとも RDMA 対応の EFA ネットワークを備えた 1 つの Prefill ノードと 1 つの Decode ノードが必要です。サポートされているインスタンスタイプについては、「HyperPod クラスターに DPD モデルエンドポイントのデプロイ」セクションをご覧ください。

アーキテクチャ

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

インテリジェントルーター

ルーターは制御プレーンです。各プロンプトをトークン化し、設定可能なトークン閾値を適用して、リクエストが分離パスを経由するか、デコーダー上でエンドツーエンドで実行されるかを決定します。長文コンテキストのプロンプトは、まず Prefiller を経由してから Decoder へ送られます。短いプロンプトは Prefiller をスキップし、有益ではないクロス GPU の KV キャッシュ転送を回避します。分離リクエストの場合、Prefiller に計算と KV キャッシュのデコーダーへのプッシュを指示し、その後生成のためにそのデコーダーへリクエストを転送します。また、replica 間でのキャッシュ局所性を最大化するために、intelligentRoutingSpec.routingStrategy を介して per-prefiller ルーティング戦略（prefixaware, kvaware, session, roundrobin）もサポートしています。

Prefiller pod

Prefiller は、LMCacheConnectorV1 を介して LMCache を KV コネクタとして使用する vLLM ワーカーです。長いプロンプトに対する KV キャッシュを計算し、LMCache の PD 送信者バックエンドレイヤーを通じて層ごとに選択されたデコーダへプッシュします。これにより、計算と転送が重なることで GPU を飽和状態に保ちます。また、LMCache は各 Prefiller に L1 CPU キャッシュも提供します。プレフィックスが再出現する場合（システムプロンプト、多回対話の履歴、検索コンテキストなど）、GPU の再計算を行わずに CPU メモリからサービスを提供するため、TTFT（Time To First Token）の大幅な向上を実現します。InferenceEndpointConfig で DPD を有効化すると、コネクタとキャッシュが自動的にプロビジョニングされます。

Decoder pod

デコーダは、LMCache を受信側として使用する vLLM ワーカーです。着信する KV 転送のために GPU メモリ（PD バッファ）を確保します。このバッファサイズは PD_BUFFER_SIZE によって決定されます。デコードカーネルには完全な CUDA グラフを実行し、転送が完了した瞬間に生成を開始します。デコーダはプレフィルを実行しないため、並列処理下でもデコードレイテンシは安定しており、長文コンテキストのリクエストを追加しても、すでにストリーミング中のトークンを妨げることはありません。

KV transfer

KV キャッシュ転送は、HyperPod がエンドツーエンドで構築する 4 レイヤーのスタック（LMCache PD → NIXL → libfabric → EFA）を使用します。LMCache の PD バックエンドが、prefiller 側のプット処理と decoder 側の取得処理をオーケストレーションします。NIXL は GPU、CPU、およびリモートピア間における統一されたメモリ抽象化を提供し、適切な RDMA オペレーションを選択します。libfabric プロバイダは EFA をカーネルバイパスかつ GPU-Direct RDMA として公開し、ホスト CPU をデータパスから除外します。これにより、転送コストは prefill 計算に対して無視できるほど小さくなります。具体的には、3,200 Gbps の EFA を備えた ml.p5.48xlarge インスタンスにおいて、Llama 3.3 70B の 8,000 トークンの転送は単一桁のミリ秒で完了します。HyperPod はこのスタックを事前統合した状態で提供するため、DPD サポート対応のワーカーイメージを選択するだけで、オペレーターが各ポッド上でコネクタ、NIXL、および EFA の接続を自動的に設定します。

Deployment overview

Disaggregated Prefill and Decode (DPD) は、HyperPod Inference Operator によって実装される機能です。本セクションでは、前提条件、推論オペレーターのインストール方法、および DPD を活用して Llama 70B モデルを効率的にサーブするための推論エンドポイントのデプロイ手順について解説します。

前提条件と HyperPod Inference Operator のインストール

AWS Command Line Interface (AWS CLI)、HyperPod クラスターへの kubectl アクセス、HuggingFace トークン、および十分なサービスクォータがあることを確認してください。ローカルの kubectl 設定を設定して、HyperPod クラスターに接続できるようにします。詳細については、Disaggregated Prefill and Decode for HyperPod inference を参照してください。

DPD (Disaggregated Prefill and Decode: 分離型プリフィルおよびデコード) には、HyperPod Inference Operator のバージョン 3.2 以降が必要です。このオペレーターは、新しい HyperPod EKS クラスターではデフォルトでインストールされています。インストール、セットアップ、アップグレードの手順については、Unlock efficient model deployment: simplified Inference Operator setup on Amazon SageMaker HyperPod を参照してください。

以下のコマンドを実行して、オペレーターのバージョンを確認してください:

kubectl get deployment hyperpod-inference-operator-controller-manager \

-n hyperpod-inference-system \

-o jsonpath='{.spec.template.spec.containers[?(@.name=="manager")].image}{"\n"}'

出力には、完全なコンテナイメージの参照が表示されます。末尾のタグにバージョン情報が含まれており、例えば以下のようになります:

XXXXXXXXXXXXXXX.dkr.ecr.us-east-2.amazonaws.com/hyperpod-inference-operator:v3.2

オペレーターバージョンが最新でない場合は、HyperPod Inference Operator リリースノート のアップグレード手順に従って、続行する前にアップグレードしてください。

HyperPod クラスターに DPD モデルエンドポイントをデプロイする

この例では、Meta Llama 3.3 70B モデルを 2 つの ml.p5.48xlarge インスタンス上でデプロイします。続行する前に、HyperPod クラスター内のインスタンスグループにこれらのインスタンスが利用可能であることを確認してください。DPD（Disaggregated Prefill and Decode）推論デプロイメントでは、NVLink と EFA（Elastic Fabric Adapter）の両方をサポートするインスタンスタイプを選択してください。EFA は読み取りおよび書き込みモードでの RDMA（Remote Direct Memory Access）をサポートしている必要があります。これには AWS 上の P5 および P6 インスタンスタイプファミリーが含まれます。なお、EFA の高帯域幅通信のためには、インスタンスは同じアベイラビリティゾーン（AZ）内に配置されている必要があります。G6、G6e、G7e インスタンスタイプファミリーも RDMA 読み取り/書き込みを備えた EFA をサポートしていますが、マルチ GPU インスタンスにおけるパフォーマンスは、PCI を介した GPU から GPU への通信によってボトルネックになります。

推論デプロイメント用のワーカーイメージには、vLLM、LMCache、NVIDIA NIXL、および EFA libfabric プロバイダーが含まれている必要があります。執筆時点では、2 つのイメージオプションをサポートしています：

• オープンソース LMCache: lmcache/vllm-openai v0.4.3。

• SageMaker Deep Learning Container (DLC): vllm:server-hyperpod-cuda-v1.1。

モデルチェックポイントの場所

HyperPod Inference Operator は、Amazon Simple Storage Service (Amazon S3) バケット、Amazon FSx ファイルシステム、HuggingFace からの直接プル、および インスタンスの NVMe ストレージ を含む、幅広いチェックポイント読み込みソースをサポートしています。本記事では、モデルチェックポイントを Amazon S3 バケットからロードします。

HyperPod クラスターと同じリージョンにある S3 バケットに、お好みのモデルチェックポイントがダウンロードされていることを確認してください。まだ行っていない場合は、バケット名と HuggingFace トークンを設定し、HuggingFace から Meta Llama 3.3 70B Instruct をダウンロードして Amazon S3 に同期してください。Amazon S3 への高帯域幅ネットワークを実現するためには、Amazon Elastic Compute Cloud (Amazon EC2) インスタンスからこの操作を実行することを推奨します。

export MODEL_BUCKET=
export MODEL_PREFIX=Llama-3.3-70B-Instruct
export AWS_REGION=
export HF_TOKEN=
pip install -U "huggingface_hub[cli]" "huggingface_hub[hf-transfer]"
HF_HUB_ENABLE_HF_TRANSFER=1 hf download meta-llama/Llama-3.3-70B-Instruct \
  --local-dir ./$MODEL_PREFIX \
  --token "$HF_TOKEN"
aws s3 sync ./$MODEL_PREFIX \
  s3://$MODEL_BUCKET/$MODEL_PREFIX/ \
  --region "$AWS_REGION"

モデルデプロイマニフェストを準備し、必要に応じて環境変数を変更してください。

export DEPLOYMENT_NAME="dpd-test-deployment"

export ENDPOINT_NAME="dpd-test"

export MODEL_NAME="meta-llama-3-3-70b"

export NAMESPACE="default"

export INSTANCE_TYPE="ml.p5.48xlarge"

export GPUS_PER_NODE="8"

export MODEL_IMAGE="lmcache/vllm-openai:v0.4.3"

完全なデプロイ YAML は、DPD エンドポイントのデプロイを参照してください。

DPD 関連のデプロイマニフェスト内のフィールド

InferenceEndpointConfig のほとんどのフィールドは、非 DPD エンドポイントと共有されており、Inference Operator ドキュメントで文書化されています。以下のフィールドは、DPD において必須であるか、または異なる意味を持ちます。

spec.pdSpec: プリフィル/デコードのトポロジを宣言し、引数を指定します。このフィールドが存在することが、エンドポイントを非集約型にする要因となります。オペレーターは、プリフィルとデコードのために個別の Deployment オブジェクトを作成し、ルーターおよび LMCache PD バックエンドを通じてこれらを接続します。

• replicas: プリフィルとデコードを独立してスケールします。

• resources: ロールのポッド仕様に適用されます。DPD ポッドではトップレベルの worker.resources は無視され、ロールごとの値が優先されます。

• routingThreshold: リクエストが非集約パスを使用するトークン長のしきい値です。このしきい値以下の場合、リクエストはプリフィルラーをバイパスして直接デコーダーへ送られます。

• args: そのロール固有の vLLM フラグで、起動時に worker.args にマージされます。worker.args に既に存在するフラグはロールごとの値に置き換えられ、存在しないフラグは追加されます。

spec.worker.environmentVariables: これらの環境変数はプリフィルラーとデコーダーのコンテナの両方に同一の方法で適用されます。現時点ではロールごとの環境変数フィールドはありません。ロールごとの動作については、代わりに pdSpec.{prefillSpec,decodingSpec}.args を使用してください。

環境変数に関する詳細は DPD エンドポイントのデプロイ をご覧ください。

マニフェストを適用してデプロイメントを検証する

kubectl apply -f inference_endpoint_dpd_config.yaml

オペレーターは、名前空間内に 2 つの Deployment オブジェクトと hyperpod-inference-system にルーター用の Deployment を作成します。イメージのプルとモデルの読み込みには数分かかります。ポッドはまず ContainerCreating 状態に入り、その後コンテナが起動すると Running 状態になります。両方の名前空間にわたってポッドを一覧表示してください：

kubectl get pods -A \

| grep -E "prefill-${DEPLOYMENT_NAME}|decode-${DEPLOYMENT_NAME}|${DEPLOYMENT_NAME}-${NAMESPACE}-router"

NAMESPACE NAME READY STATUS RESTARTS AGE

default prefill-dpd-test-deployment-XXXX 3/3 Running 0 7m

default decode-dpd-test-deployment-XXXX 3/3 Running 0 7m

hyperpod-inference-system dpd-test-deployment-default-router-XXXX 2/2 Running 0 7m

各モデルポッドには 3 つのコンテナが含まれています：vLLM ワーカー、Nginx リバースプロキシ（reverse proxy）、および OpenTelemetry カレータです。ルーターポッドには 2 つのコンテナ（router, otel）があります。IEC 条件は準備完了を報告します：

kubectl get inferenceendpointconfig ${DEPLOYMENT_NAME} -n ${NAMESPACE} \

-o jsonpath='{.status.conditions[0].message}{"\n"}'

DPD prefill および decode デプロイメントが準備完了しました

エンドポイントの呼び出しと KV 転送の確認

エンドポイントが準備できたら、短いプロンプトと長いプロンプトを送信して両方のルーティングパスを動作させます。prefiller と decoder のログを確認し、KV キャッシュ（key-value cache）が EFA を介して転送されていることを確認します。以下のコマンドは、IEC が前述のマニフェストと同じ ${NAMESPACE}、${ENDPOINT_NAME}、および ${DEPLOYMENT_NAME} でデプロイされたことを前提としています。

必要なポッド名とルーター URL を取得します：

PREFILL_POD=$(kubectl get pod -n ${NAMESPACE} \

-l 'inference.sagemaker.aws.amazon.com/dpd-role=prefill' \

-o jsonpath='{.items[0].metadata.name}')

DECODE_POD=$(kubectl get pod -n ${NAMESPACE} \

-l 'inference.sagemaker.aws.amazon.com/dpd-role=decode' \

-o jsonpath='{.items[0].metadata.name}')

ROUTER_POD=$(kubectl get pods -n hyperpod-inference-system -o name \

| grep -- "${DEPLOYMENT_NAME}-${NAMESPACE}-router" | head -1)

ROUTER_URL=http://${DEPLOYMENT_NAME}-${NAMESPACE}-routing-service.hyperpod-inference-system.svc.cluster.local:443/v1/chat/completions

ショートプロンプト（閾値未満、デコーダへの直接パス）

クラスター内のポッドでエンドポイントを呼び出すショートプロンプトを実行します:

kubectl run curl-short --rm -it --image=curlimages/curl --restart=Never -- \
  curl -s -k -X POST "$ROUTER_URL" \
    -H "Content-Type: application/json" \
    -d '{
      "model": "/opt/ml/model",
      "messages": [{"role": "user", "content": "What is disaggregated prefill-decode in one sentence?"}],
      "max_tokens": 80,
      "temperature": 0.0
    }'

{

"id": "chatcmpl-7-...",

"object": "chat.completion",

"model": "/opt/ml/model",

"choices": [{

"index": 0,

"message": {"role": "assistant", "content": "Disaggregated prefill-decode is an inference architecture that ..."},

"finish_reason": "stop"

}]

}

Long prompt (above threshold, DPD path)

4,096 トクンのルーティング閾値を超えるプロンプトは、まずプリフェッチャ（prefiller）を経由し、その後トークン生成のためにデコーダーへルーティングされます。以下の例では、ある文を繰り返し追加して約 6,000 トクンのプロンプトを構築しています:

kubectl run curl-long --rm -it --image=curlimages/curl --restart=Never -- sh -c '

ROUTER="'$ROUTER_URL'"

LONG=""

i=0; while [ $i -lt 600 ]; do LONG="${LONG}The quick brown