View a markdown version of this page

階層型 KV キャッシュの設定 - Amazon SageMaker AI
前提条件RayCluster を設定するRay Serve アプリケーションを設定するCPU オフロードを有効にする (オプション)

翻訳は機械翻訳により提供されています。提供された翻訳内容と英語版の間で齟齬、不一致または矛盾がある場合、英語版が優先します。

階層型 KV キャッシュの設定

マネージド階層型 KV キャッシュは、Ray Serve デプロイにクラスター全体のキャッシュ階層を追加するため、あるレプリカは別のレプリカがすでに計算したトークンプレフィックスを読み取ります。これにより、長いドキュメント、マルチターン会話、共有システムプロンプトの最初のトークンまでの時間が短縮されます。クラスター全体の階層は、LMCache ドキュメントの HyperPod ストレージバックエンドと LMCache を使用して HyperPod 階層型ストレージで実行されます。これは Ray Serve および vLLM と統合され、HyperPod 階層型ストレージにバックアップされた分散 KV キャッシュを提供します。 LMCache HyperPod

キャッシュには 2 つの階層があります。

  • 最も迅速に再利用できるように、リクエストを処理するノードのメモリ内のローカル階層。

  • HyperPod 階層型ストレージのクラスター全体の階層。クラスターノードにまたがるプールされたメモリ階層。レプリカは、再計算するのではなく、別のレプリカによって計算されたプレフィックスを読み取ります。

リクエストがトークンプレフィックスを以前の作業と共有すると、デプロイはキャッシュされた KV 状態をローカル階層、次にクラスター全体の階層から読み取り、その後に何かを再計算します。

前提条件

  • Amazon EKS によってオーケストレーションされた HyperPod クラスター。KubeRay 演算子がインストールされています。

  • クラスターで HyperPod 階層型ストレージが有効になっています。セットアップ手順については、「マネージド階層型チェックポイントの設定」を参照してください。

RayCluster を設定する

ワーカーグループの仕様に以下を追加して、階層型ストレージエンドポイントと共有メモリを Ray Serve レプリカに公開します。

workerGroupSpecs:
- template:
    spec:
      volumes:
        - name: host-shm
          hostPath:
            path: /dev/shm
      containers:
      - name: ray-worker
        env:
          - name: NODE_IP
            valueFrom:
              fieldRef:
                fieldPath: status.hostIP
        volumeMounts:
          - name: host-shm
            mountPath: /dev/shm/ai_toolkit_cache
            subPath: ai_toolkit_cache

NODE_IP 環境変数はホスト IP アドレスを提供するため、LMCache クライアントはポート 9200 で階層型ストレージサービスに接続できます。共有メモリボリュームマウントは、ローカル KV キャッシュの CPU オフロードパスを有効にします。

Ray Serve アプリケーションを設定する

LMCache コネクタで LLMConfig API を使用して、Ray Serve デプロイで KV キャッシュを有効にします。次の例では、階層型 KV キャッシュが有効になっているモデルを提供します。

from ray.serve.llm import LLMConfig, build_openai_app

llm_config = LLMConfig(
    model_loading_config={
        "model_id": "my-llm",
        "model_source": "Qwen/Qwen-7B-Chat",
    },
    engine_kwargs={
        "trust_remote_code": True,
        "kv_transfer_config": {
            "kv_connector": "LMCacheConnectorV1",
            "kv_role": "kv_both",
        },
    },
    runtime_env={
        "env_vars": {
            "LMCACHE_REMOTE_URL": "sagemaker-hyperpod://$(NODE_IP):9200",
            "LMCACHE_EXTRA_CONFIG": '{"sagemaker_hyperpod_bucket": "lmcache", "sagemaker_hyperpod_shared_memory_name": "ai_toolkit_cache"}',
            "LMCACHE_CHUNK_SIZE": "256",
        },
    },
    accelerator_type="A10G",
)

app = build_openai_app({"llm_configs": [llm_config]})

キー設定値は次のとおりです。

  • kv_connector: に設定するとLMCacheConnectorV1、LMCache と vLLM の統合が有効になります。

  • kv_role: 各レプリカが共有キャッシュに対して読み取りと書き込みの両方kv_bothを行うように を設定します。

  • LMCACHE_REMOTE_URL: ローカルノードの階層型ストレージエンドポイント。RayCluster マニフェストで設定されたNODE_IP環境変数を使用します。

  • LMCACHE_CHUNK_SIZE: キャッシュチャンクあたりのトークンの数。値が小さいほど、より多くのキャッシュエントリを犠牲にして、共有プレフィックスのキャッシュヒット率が向上します。

CPU オフロードを有効にする (オプション)

階層型ストレージに書き込まれる前に KV エントリをキャッシュするローカル CPU メモリ層を追加するには、次の環境変数を に追加しますruntime_env。

"env_vars": {
    "LMCACHE_REMOTE_URL": "sagemaker-hyperpod://$(NODE_IP):9200",
    "LMCACHE_EXTRA_CONFIG": '{"sagemaker_hyperpod_bucket": "lmcache", "sagemaker_hyperpod_shared_memory_name": "ai_toolkit_cache"}',
    "LMCACHE_CHUNK_SIZE": "256",
    "LMCACHE_LOCAL_CPU": "True",
    "LMCACHE_MAX_LOCAL_CPU_SIZE": "100",
}

CPU オフロードが有効になっている場合、KV キャッシュエントリは、クラスター全体の階層型ストレージに書き込まれる前に、ローカルノードの CPU メモリに保存されます。これにより、同じノード上のレプリカのキャッシュ読み取りが高速化されます。