Ray と TorchTPU を使用してマルチホスト TPU で LLM をサービングする

このチュートリアルでは、Ray Serve LLM、vLLM、vLLM-TorchTPU を使用して、Google Kubernetes Engine(GKE)のマルチホスト Cloud TPU スライスでテンソル並列処理とエキスパート並列処理を使用して、Mixture-of-Experts(MoE)Gemma 4 26B-A4B-it 大規模言語モデル(LLM)をサービングする方法について説明します。

このチュートリアルは、Kubernetes コンテナ オーケストレーションを使用して、分散マルチホスト TPU スライスで AI / ML ワークロードを提供する機械学習(ML)エンジニア、プラットフォーム管理者、オペレーター、データおよび AI スペシャリストを対象としています。 Google Cloud のコンテンツで使用されている一般的なロールとタスクの例の詳細については、一般的な GKE ユーザーロールとタスクをご覧ください。

このチュートリアルを開始する前に、次の内容をよく理解しておいてください。

背景

このセクションでは、このチュートリアルで使用されている重要なテクノロジーについて説明します。

TPU

Tensor Processing Unit(TPU)は、機械学習ワークロード用に Google がカスタム設計したアクセラレータです。このチュートリアルは、TPU Trillium(v6e)と Ironwood(TPU7x)の両方をサポートしています。マルチホスト TPU スライスでは、複数の TPU ホストが高速チップ間相互接続(ICI)ネットワークを介して直接通信し、単一のチップまたは TPU ホストのメモリ容量を超えるモデルの提供に必要な、集約された高帯域幅メモリ(HBM)と低レイテンシの集団通信を提供します。

vLLM と TorchTPU

vLLM は、LLM の推論およびサービングを行うためのオープンソース エンジンです。TorchTPU は、Cloud TPU 用のオープンソースの PyTorch バックエンドです。ICI を介したデバイス管理と集団通信を提供します。

vLLM-TorchTPU は、TPU プラットフォーム プラグインを vLLM に登録し、TorchTPU と Accelerated Linear Algebra(XLA)コンパイラを介して Cloud TPU の PyTorch モデルのフォワードパスをコンパイルします。Ray を使用して、マルチホスト TPU スライスに分散されたテンソル並列処理とエキスパート並列処理をサポートします。

GKE での Ray Serve LLM

Ray Serve LLM は、/v1/chat/completions、/v1/completions、/v1/models エンドポイントに OpenAI 互換の HTTP Ingress を提供し、GKE で分散 vLLM デプロイをオーケストレートします。Ray on GKE アドオンを有効にすると、KubeRay オペレーターとその TPU Webhook が、スライス トポロジ メタデータ、ワーカー ホスト名、TorchTPU 調整変数を各マルチホスト TPU ワーカー Pod に挿入します。

Ray Serve LLM には、マルチホスト TPU サービング用に次の機能が用意されています。

  • トポロジを考慮したスライス配置では、accelerator_config フィールドに TPU トポロジを指定すると、各モデルレプリカに完全なマルチホスト TPU スライスが予約され、スライス全体で TPU デバイスごとに 1 つの vLLM ワーカーがスケジュールされます。
  • プレフィックス対応ルーティングは、複数の TPU スライス レプリカにわたってサービングする際に、共有システム プロンプトまたはマルチターン会話履歴を含むリクエストを、TPU HBM にそれらのプレフィックス Key-Value(KV)ブロックをすでにキャッシュしているレプリカに転送します。
  • 直接トークン ストリーミングは、同時リクエスト間でアウトバウンド トークン ストリームをバッチ処理し、モデルレプリカからレスポンスを直接ストリーミングして、最初のトークンまでの時間(TTFT)とトークン間レイテンシ(ITL)を短縮します。
  • 同時実行ベースのレプリカ自動スケーリングは、進行中のリクエストの同時実行数に基づいて、モデルレプリカをスライス全体単位でスケーリングします。

目標

このチュートリアルでは、次の方法を説明します。

  1. 事前ビルドされた vLLM-TorchTPU ベースイメージの上に Ray Serve をレイヤ化するコンテナ イメージをビルドし、Artifact Registry に push します。
  2. TPU Trillium(v6e)または Ironwood(TPU7x)で Autopilot または Standard の GKE クラスタを作成します。
  3. マルチホスト RayCluster カスタム リソースをデプロイし、TPU トポロジ構成を確認します。
  4. Ray Serve LLM アプリケーションをデプロイして、マルチホスト TPU スライス全体でテンソル並列処理とエキスパート並列処理を使用して google/gemma-4-26B-A4B-it モデルを提供します。または、RayService カスタム リソースを使用してクラスタとアプリケーションの両方を宣言的にデプロイします。
  5. OpenAI 互換のモデル検出リクエストとチャット補完リクエストを、サービングされたモデルに送信します。

費用

のドキュメントでは、課金対象である次の Google Cloudコンポーネントを使用します。

料金計算ツールを使うと、予想使用量に基づいて費用の見積もりを算出できます。

このチュートリアルに記載されているタスクの完了後、作成したリソースを削除すると、それ以上の請求は発生しません。詳細については、クリーンアップをご覧ください。

始める前に

作業を始める前に、次のタスクが完了していることを確認してください。

  • Artifact Registry API と Google Kubernetes Engine API を有効にします。
  • API を有効にする
  • このタスクに Google Cloud CLI を使用する場合は、gcloud CLI をインストールして初期化します。gcloud CLI をインストール済みの場合は、gcloud components update コマンドを実行して最新のバージョンを取得します。以前のバージョンの gcloud CLI では、このドキュメントのコマンドを実行できない場合があります。

環境を準備する

このチュートリアルでは、Cloud Shell またはローカル ワークステーションを使用して、 Google Cloudでホストされているリソースを管理します。Cloud Shell には、このチュートリアルで使用するコマンドライン ツール(kubectl、docker、python3、gcloud CLI など)がプリインストールされています。

Cloud Shell を使用して環境を設定するには、次の操作を行います。

  1. Google Cloud コンソールで Cloud Shell 有効化アイコン。(Cloud Shell をアクティブにする)をクリックして、Google Cloud コンソールで Cloud Shell セッションを起動します。このアクションにより、 Google Cloud コンソールの下部ペインでセッションが起動します。

  2. kubernetes-engine-samples リポジトリのクローンを作成し、このチュートリアルのサンプル ディレクトリに移動します。

    git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples.git
    cd kubernetes-engine-samples/ai-ml/gke-ray/rayserve/llm/torchtpu
    
  3. Ray CLI がインストールされたローカル Python 仮想環境を作成して有効にします。

    python3 -m venv ~/ray-env
    source ~/ray-env/bin/activate
    pip install "ray[default]==2.58.0"
    
  4. Google Cloud プロジェクトとクラスタの環境変数を設定します。

    export PROJECT_ID=PROJECT_ID
    export CLUSTER_NAME=CLUSTER_NAME
    export REGION=REGION
    export ZONE=ZONE
    export ARTIFACT_REGISTRY=ARTIFACT_REGISTRY
    export HF_TOKEN=HF_TOKEN
    export IMAGE_URI=${REGION}-docker.pkg.dev/${PROJECT_ID}/${ARTIFACT_REGISTRY}/ray-torchtpu-vllm:latest
    gcloud config set project ${PROJECT_ID}
    

    次のように置き換えます。

    • PROJECT_ID: 実際の Google Cloud プロジェクト ID。
    • CLUSTER_NAME: GKE クラスタの名前(ray-serve-torchtpu-cluster など)。
    • REGION: Artifact Registry リポジトリと TPU 容量が配置されているリージョン(us-east5 など)。
    • ZONE: TPU の容量が使用可能なゾーン(us-east5-b など)。詳細については、GKE で TPU の可用性を検証するをご覧ください。
    • ARTIFACT_REGISTRY: Artifact Registry Docker リポジトリの名前(例: ray-torchtpu)。
    • HF_TOKEN: Read 権限を持つ Hugging Face アクセス トークン。

カスタム コンテナ イメージをビルドする

TorchTPU、vLLM、vLLM-TorchTPU プラグインを含む、ビルド済みの vLLM-TorchTPU ベースイメージの上に ray[default,serve] パッケージをインストールするコンテナ イメージをビルドします。

  1. kubernetes-engine-samples/ai-ml/gke-ray/rayserve/llm/torchtpu ディレクトリの Dockerfile を確認します。

    FROM vllm/vllm-torchtpu:v0.29.0
    
    ENV HF_HOME=/dev/shm/huggingface
    ENV RAY_SERVE_DEFAULT_HTTP_HOST=0.0.0.0
    
    # Install wget for KubeRay health probes and pin Ray Serve and transformers
    # versions compatible with vLLM 0.29.0 and multi-host TPUs.
    RUN apt-get update && apt-get install -y --no-install-recommends wget \
        && rm -rf /var/lib/apt/lists/* \
        && pip install --no-cache-dir \
           "ray[default,serve] @ https://s3-us-west-2.amazonaws.com/ray-wheels/master/6b2b414bb7a5b4e2705ed9f4fdab9787b92f3b08/ray-3.0.0.dev0-cp312-cp312-manylinux2014_x86_64.whl" \
           "transformers>=5.18.0"

    Ray Serve LLM を使用したマルチホスト TPU サービングには、TPU Trillium(v6e)では Ray バージョン 2.56.0 以降、Ironwood(TPU7x)でのチップレットごとのスケジューリングには Ray バージョン 2.58.0 以降が必要です。

  2. Artifact Registry Docker リポジトリを作成し、コンテナ イメージをビルドして Artifact Registry に push します。

    gcloud artifacts repositories create ${ARTIFACT_REGISTRY} \
        --repository-format=docker \
        --location=${REGION} \
        --description="Ray and vLLM-TorchTPU container images"
    
    gcloud auth configure-docker ${REGION}-docker.pkg.dev,us-docker.pkg.dev
    docker build -t ${IMAGE_URI} .
    docker push ${IMAGE_URI}
    

GKE クラスタを作成する

マネージド Ray Operator アドオンが有効になっている GKE Autopilot クラスタまたは Standard クラスタのマルチホスト TPU スライスで LLM を提供できます。

  1. GKE クラスタを作成し、kubectl コマンドライン ツールの認証情報を構成します。

    Autopilot

    1. Ray Operator アドオンを有効にして、リージョン GKE Autopilot クラスタを作成します。

      gcloud container clusters create-auto ${CLUSTER_NAME} \
          --project=${PROJECT_ID} \
          --region=${REGION} \
          --cluster-version=1.36 \
          --enable-ray-operator
      

      後の手順で RayCluster または RayService カスタム リソースをデプロイすると、GKE Autopilot は、Pod リソース リクエストとノードセレクタに一致する CPU ヘッドノードとマルチホスト TPU スライスを自動的にプロビジョニングします。

    2. クラスタと通信するように kubectl コマンドライン ツールを構成します。

      gcloud container clusters get-credentials ${CLUSTER_NAME} \
          --project=${PROJECT_ID} \
          --region=${REGION}
      

    標準

    1. Ray ヘッド Pod をホストする CPU ノードプール、イメージ ストリーミング、GKE Dataplane V2 を有効にして、ゾーン GKE Standard クラスタを作成します。

      gcloud container clusters create ${CLUSTER_NAME} \
          --project=${PROJECT_ID} \
          --zone=${ZONE} \
          --cluster-version=1.36 \
          --machine-type=n2-standard-8 \
          --num-nodes=1 \
          --addons=RayOperator \
          --enable-image-streaming \
          --enable-dataplane-v2
      
    2. クラスタと通信するように kubectl コマンドライン ツールを構成します。

      gcloud container clusters get-credentials ${CLUSTER_NAME} \
          --project=${PROJECT_ID} \
          --zone=${ZONE}
      
    3. TPU Trillium(v6e)の 4 つのホストにまたがる 4x4 トポロジ、または Ironwood(TPU7x)の 2 つのホストにまたがる 2x2x2 トポロジを使用して、16 デバイスのマルチホスト TPU スライス ノードプールを作成します。

      Trillium(v6e)

      gcloud container node-pools create tpu-multi-host-pool \
          --project=${PROJECT_ID} \
          --cluster=${CLUSTER_NAME} \
          --zone=${ZONE} \
          --machine-type=ct6e-standard-4t \
          --tpu-topology=4x4 \
          --num-nodes=4
      

      Ironwood(TPU7x)

      gcloud container node-pools create tpu-multi-host-pool \
          --project=${PROJECT_ID} \
          --cluster=${CLUSTER_NAME} \
          --zone=${ZONE} \
          --machine-type=tpu7x-standard-4t \
          --tpu-topology=2x2x2 \
          --num-nodes=2
      
  2. Ray Pod がモデルの重みをダウンロードするときに Hugging Face で認証できるように、Hugging Face トークンを使用して Kubernetes Secret を作成します。

    kubectl create secret generic hf-secret \
        --from-literal=hf_api_token=${HF_TOKEN}
    

マルチホスト RayCluster カスタム リソースをデプロイする

CPU ヘッド Pod、Ray 共有メモリ オブジェクト ストア用のインメモリ /dev/shm ボリューム、マルチホスト TPU ワーカー グループをプロビジョニングする RayCluster カスタム リソースをデプロイします。

  1. サンプル リポジトリで、TPU バージョンの RayCluster マニフェストを確認します。

    Trillium(v6e)

    apiVersion: ray.io/v1
    kind: RayCluster
    metadata:
      name: ray-serve-torchtpu
    spec:
      headGroupSpec:
        rayStartParams:
          num-cpus: "0"
        template:
          spec:
            containers:
            - name: ray-head
              image: ${IMAGE_URI}
              env:
              - name: HF_TOKEN
                valueFrom:
                  secretKeyRef:
                    name: hf-secret
                    key: hf_api_token
              ports:
              - containerPort: 6379
                name: gcs
              - containerPort: 8265
                name: dashboard
              - containerPort: 10001
                name: client
              - containerPort: 8000
                name: serve
              resources:
                limits:
                  cpu: "4"
                  memory: "16Gi"
                requests:
                  cpu: "4"
                  memory: "16Gi"
      workerGroupSpecs:
      - groupName: v6e-16
        replicas: 1
        numOfHosts: 4
        template:
          spec:
            containers:
            - name: ray-worker
              image: ${IMAGE_URI}
              env:
              - name: HF_TOKEN
                valueFrom:
                  secretKeyRef:
                    name: hf-secret
                    key: hf_api_token
              resources:
                limits:
                  cpu: "64"
                  memory: "200Gi"
                  google.com/tpu: "4"
                requests:
                  cpu: "64"
                  memory: "200Gi"
                  google.com/tpu: "4"
              volumeMounts:
              - mountPath: /dev/shm
                name: dshm
            volumes:
            - name: dshm
              emptyDir:
                medium: Memory
                sizeLimit: 100Gi
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
              cloud.google.com/gke-tpu-topology: 4x4

    Ironwood(TPU7x)

    apiVersion: ray.io/v1
    kind: RayCluster
    metadata:
      name: ray-serve-torchtpu
    spec:
      headGroupSpec:
        rayStartParams:
          num-cpus: "0"
        template:
          spec:
            containers:
            - name: ray-head
              image: ${IMAGE_URI}
              env:
              - name: HF_TOKEN
                valueFrom:
                  secretKeyRef:
                    name: hf-secret
                    key: hf_api_token
              - name: RAY_TPU_RESOURCE_PER_CHIP
                value: "2"
              ports:
              - containerPort: 6379
                name: gcs
              - containerPort: 8265
                name: dashboard
              - containerPort: 10001
                name: client
              - containerPort: 8000
                name: serve
              resources:
                limits:
                  cpu: "4"
                  memory: "16Gi"
                requests:
                  cpu: "4"
                  memory: "16Gi"
      workerGroupSpecs:
      - groupName: tpu7x-16
        replicas: 1
        numOfHosts: 2
        resources:
          # Advertises 8 logical TPU cores (2 chiplets per physical chip) per host.
          TPU: "8"
        template:
          spec:
            containers:
            - name: ray-worker
              image: ${IMAGE_URI}
              env:
              - name: HF_TOKEN
                valueFrom:
                  secretKeyRef:
                    name: hf-secret
                    key: hf_api_token
              - name: RAY_TPU_RESOURCE_PER_CHIP
                value: "2"
              resources:
                limits:
                  cpu: "64"
                  memory: "250Gi"
                  google.com/tpu: "4"
                requests:
                  cpu: "64"
                  memory: "250Gi"
                  google.com/tpu: "4"
              volumeMounts:
              - mountPath: /dev/shm
                name: dshm
            volumes:
            - name: dshm
              emptyDir:
                medium: Memory
                sizeLimit: 128Gi
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu7x
              cloud.google.com/gke-tpu-topology: 2x2x2

    Ironwood(TPU7x)では、各物理チップに 2 つのチップレットが含まれており、4 チップホストあたり 8 つの論理 TPU コアが提供されます。このマニフェストでは、resources セクションで TPU フィールドの値を "8" に設定し、RAY_TPU_RESOURCE_PER_CHIP 環境変数を "2" に設定して、ホストごとに 8 つの論理 TPU リソースをアドバタイズしてスケジュールします。

  2. マニフェストを適用して RayCluster カスタム リソースを作成します。

    Trillium(v6e)

    envsubst < ray-cluster.tpu-v6e-multihost.yaml | kubectl apply -f -
    

    Ironwood(TPU7x)

    envsubst < ray-cluster.tpu-v7x-multihost.yaml | kubectl apply -f -
    
  3. ヘッド Pod とすべての TPU ワーカー Pod が Running 状態と 1/1 READY 状態になるまで待ちます。

    kubectl get pods -l ray.io/cluster=ray-serve-torchtpu
    

    出力は次のようになります。

    Trillium(v6e)

    NAME                                     READY   STATUS    RESTARTS   AGE
    ray-serve-torchtpu-head-5cj7g            1/1     Running   0          35s
    ray-serve-torchtpu-v6e-16-worker-6x8cg   1/1     Running   0          35s
    ray-serve-torchtpu-v6e-16-worker-79j5l   1/1     Running   0          35s
    ray-serve-torchtpu-v6e-16-worker-psxbl   1/1     Running   0          35s
    ray-serve-torchtpu-v6e-16-worker-rz7dv   1/1     Running   0          35s
    

    Ironwood(TPU7x)

    NAME                                       READY   STATUS    RESTARTS   AGE
    ray-serve-torchtpu-head-w7b2v              1/1     Running   0          2m30s
    ray-serve-torchtpu-tpu7x-16-worker-6zlnj   1/1     Running   0          2m30s
    ray-serve-torchtpu-tpu7x-16-worker-cb4zz   1/1     Running   0          2m30s
    
  4. KubeRay TPU Webhook がマルチホスト スライス、ワーカー調整、TorchTPU 環境変数をワーカー Pod に挿入したことを確認します。

    kubectl get pods -l ray.io/node-type=worker \
        -o=jsonpath='{range .items[0].spec.containers[0].env[*]}{.name}={.value}{"\n"}{end}' \
        | grep -E "TPU_WORKER_ID|TPU_WORKER_HOSTNAMES|TPU_NAME|TORCH_TPU_TOPOLOGY|TORCH_TPU_SLICEBUILDER_ADDRESSES"
    

    出力は次のようになります。

    Trillium(v6e)

    TPU_WORKER_HOSTNAMES=v6e-16-0-0.ray-serve-torchtpu-headless,v6e-16-0-1.ray-serve-torchtpu-headless,v6e-16-0-2.ray-serve-torchtpu-headless,v6e-16-0-3.ray-serve-torchtpu-headless
    TORCH_TPU_TOPOLOGY=4,4,1
    TORCH_TPU_SLICEBUILDER_ADDRESSES=v6e-16-0-0.ray-serve-torchtpu-headless:8471,v6e-16-0-0.ray-serve-torchtpu-headless:8472,v6e-16-0-0.ray-serve-torchtpu-headless:8473,v6e-16-0-0.ray-serve-torchtpu-headless:8474,v6e-16-0-1.ray-serve-torchtpu-headless:8471,v6e-16-0-1.ray-serve-torchtpu-headless:8472,v6e-16-0-1.ray-serve-torchtpu-headless:8473,v6e-16-0-1.ray-serve-torchtpu-headless:8474,v6e-16-0-2.ray-serve-torchtpu-headless:8471,v6e-16-0-2.ray-serve-torchtpu-headless:8472,v6e-16-0-2.ray-serve-torchtpu-headless:8473,v6e-16-0-2.ray-serve-torchtpu-headless:8474,v6e-16-0-3.ray-serve-torchtpu-headless:8471,v6e-16-0-3.ray-serve-torchtpu-headless:8472,v6e-16-0-3.ray-serve-torchtpu-headless:8473,v6e-16-0-3.ray-serve-torchtpu-headless:8474
    TPU_WORKER_ID=1
    TPU_NAME=v6e-16-0
    

    Ironwood(TPU7x)

    TPU_WORKER_HOSTNAMES=tpu7x-16-0-0.ray-serve-torchtpu-headless,tpu7x-16-0-1.ray-serve-torchtpu-headless
    TORCH_TPU_TOPOLOGY=2,2,2,2
    TORCH_TPU_SLICEBUILDER_ADDRESSES=tpu7x-16-0-0.ray-serve-torchtpu-headless:8471,tpu7x-16-0-0.ray-serve-torchtpu-headless:8472,tpu7x-16-0-0.ray-serve-torchtpu-headless:8473,tpu7x-16-0-0.ray-serve-torchtpu-headless:8474,tpu7x-16-0-0.ray-serve-torchtpu-headless:8475,tpu7x-16-0-0.ray-serve-torchtpu-headless:8476,tpu7x-16-0-0.ray-serve-torchtpu-headless:8477,tpu7x-16-0-0.ray-serve-torchtpu-headless:8478,tpu7x-16-0-1.ray-serve-torchtpu-headless:8471,tpu7x-16-0-1.ray-serve-torchtpu-headless:8472,tpu7x-16-0-1.ray-serve-torchtpu-headless:8473,tpu7x-16-0-1.ray-serve-torchtpu-headless:8474,tpu7x-16-0-1.ray-serve-torchtpu-headless:8475,tpu7x-16-0-1.ray-serve-torchtpu-headless:8476,tpu7x-16-0-1.ray-serve-torchtpu-headless:8477,tpu7x-16-0-1.ray-serve-torchtpu-headless:8478
    TPU_WORKER_ID=1
    TPU_NAME=tpu7x-16-0
    

    TPU ワーカーノードと自動検出された ray.io/tpu-* ラベルは、Ray ダッシュボードの [Cluster] タブで確認できます。

  5. ローカルポートをポート 8265 の Ray ダッシュボードとポート 8000 の Ray Serve HTTP エンドポイントに転送します。

    kubectl port-forward svc/ray-serve-torchtpu-head-svc 8265:8265 8000:8000 >/dev/null 2>&1 &
    

Ray Serve LLM アプリケーションを作成してデプロイする

RayCluster カスタム リソースに Ray Serve LLM アプリケーションをデプロイして、マルチホスト TPU スライス全体で google/gemma-4-26B-A4B-it モデルを提供します。

  1. サンプル リポジトリの serve_vllm_torchtpu.py スクリプトを確認します。

    from ray.serve.deployment import Application
    from ray.serve.llm import (
        LLMConfig,
        LLMServingArgs,
        ModelLoadingConfig,
        build_openai_app,
    )
    
    
    def app(args: dict[str, str]) -> Application:
        llm_config = LLMConfig(
            model_loading_config=ModelLoadingConfig(
                model_id=args.get("model_id", "google/gemma-4-26B-A4B-it"),
            ),
            accelerator_type=args.get("accelerator_type", "TPU-V6E"),
            accelerator_config={
                "kind": "tpu",
                "topology": args.get("topology", "4x4"),
            },
            placement_group_config={
                "bundle_per_worker": {"TPU": 1},
            },
            engine_kwargs={
                "tensor_parallel_size": int(args.get("tensor_parallel_size", "16")),
                "enable_expert_parallel": True,
                "enable_prefix_caching": True,
                "language_model_only": True,
                "disable_hybrid_kv_cache_manager": True,
                "max_model_len": int(args.get("max_model_len", "8192")),
                "max_num_batched_tokens": int(
                    args.get("max_num_batched_tokens", "4096")
                ),
            },
            experimental_configs={
                "stream_batching_interval_ms": 10,
            },
            runtime_env={
                "env_vars": {
                    "TPU_MULTIHOST_BACKEND": "ray",
                    # Execute Gemma 4 routed expert layers on TPU TensorCores instead of SparseCore.
                    "USE_MOE_SPARSE_CORE": "0",
                    "VLLM_RAY_EXTRA_ENV_VARS_TO_COPY": "USE_MOE_SPARSE_CORE",
                }
            },
        )
        return build_openai_app(LLMServingArgs(llm_configs=[llm_config]))

    このスクリプトは、マルチホスト TPU サービング用に次の設定で LLMConfig オブジェクトを構成します。

    • accelerator_config フィールドと placement_group_config フィールドは、単一の TPU スライスにプレースメント グループを予約し、スライス全体で vLLM ワーカーごとに 1 つの TPU コアを割り当てます。
    • engine_kwargs フィールドは、テンソル並列処理とエキスパート並列処理を使用して、16 個の TPU コアすべてにモデルをシャーディングし、プレフィックス キャッシュを有効にします。また、マルチモーダル入力を無効にしてテキスト推論用に TPU HBM を予約し、Gemma 4 アテンション レイヤに統合された KV キャッシュ レイアウトを使用します。コンテキストの長さは 8,192 トークン、プリフィル バッチサイズは 4,096 トークンに制限されます。
    • experimental_configs フィールドは、送信トークン ストリームを 10 ミリ秒ごとにバッチ処理して、同時リクエスト間のストリーミング オーバーヘッドを削減します。
    • runtime_env フィールドは、vLLM-TorchTPU(TPU_MULTIHOST_BACKEND: "ray")の Ray マルチホスト バックエンドを有効にし、SparseCore ではなく TPU TensorCore で Gemma 4 MoE エキスパート レイヤを実行し、その設定を 16 個のワーカーアクターすべてに伝播します。
  2. serve run コマンドを使用して、RayCluster カスタム リソースに Ray Serve アプリケーションをデプロイします。このコマンドは、モデルの重みの読み込みが完了し、16 台の TPU デバイスすべてで XLA グラフのコンパイルが完了し、Serve アプリケーションが RUNNING 状態になるまで、デプロイログをストリーミングします。

    Trillium(v6e)

    ray job submit \
        --address http://localhost:8265 \
        --working-dir . \
        --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
        -- serve run --non-blocking serve_vllm_torchtpu:app accelerator_type=TPU-V6E topology=4x4
    

    Ironwood(TPU7x)

    ray job submit \
        --address http://localhost:8265 \
        --working-dir . \
        --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
        -- serve run --non-blocking serve_vllm_torchtpu:app accelerator_type=TPU-V7X topology=2x2x2
    

    出力は次のようになります。

    Trillium(v6e)

    (ServeController pid=1637) INFO 2026-10-13 10:23:05,160 controller 1637 -- Deploying new version of Deployment(name='LLMServer:google--gemma-4-26B-A4B-it', app='default') (initial target replicas: 1).
    (ServeController pid=1637) INFO 2026-10-13 10:23:05,161 controller 1637 -- Deploying new version of Deployment(name='OpenAiIngress', app='default') (initial target replicas: 1).
    (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) INFO 10-13 10:24:06 [__init__.py:318] TPU info: node_name=v6e-16-0 | tpu_type=v6e-16 | worker_id=1 | num_chips=4 | num_cores_per_chip=1
    (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) INFO 10-13 10:24:06 [__init__.py:238] Platform plugin torchtpu is activated
    (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) INFO 10-13 10:24:08 [core.py:112] Initializing a V1 LLM engine (v0.22.1) with config: model='google/gemma-4-26B-A4B-it', speculative_config=None, tokenizer='google/gemma-4-26B-A4B-it', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.bfloat16, max_seq_len=8192, download_dir=None, load_format=auto, tensor_parallel_size=16, pipeline_parallel_size=1, data_parallel_size=1, decode_context_parallel_size=1, dcp_comm_backend=ag_rs, disable_custom_all_reduce=True, quantization=None, quantization_config=None, enforce_eager=False, enable_return_routed_experts=False, kv_cache_dtype=auto, device_config=tpu, ...
    (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) (RayWorkerWrapper pid=1017, ip=10.72.14.21) INFO 10-13 10:24:30 [expert_map_manager.py:245] [EP Rank 4/16] Expert parallelism is enabled. Expert placement strategy: linear. Local/global number of experts: 8/128. Experts local to global index map: 0->32, 1->33, 2->34, 3->35, 4->36, 5->37, 6->38, 7->39.
    (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) (RayWorkerWrapper pid=856, ip=10.72.6.21) INFO 10-13 10:27:55 [default_loader.py:397] Loading weights took 12.76 seconds
    (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) (RayWorkerWrapper pid=856, ip=10.72.6.21) INFO 10-13 10:32:53 [tpu_worker.py:376] Memory statistics | total_hbm_limit_gb=31.25GiB | total_hbm_limit_cap_gb=28.75GiB | total_hbm_used_gb=3.69GiB | kv_cache_headroom_gb=5.0GiB | kv_connector_hbm_reserve_gb=0.0GiB | total_hbm_avail_gb=20.06GiB
    (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) INFO 10-13 10:32:54 [kv_cache_utils.py:1733] GPU KV cache size: 600,832 tokens
    (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) INFO 10-13 10:32:54 [kv_cache_utils.py:1734] Maximum concurrency for 8,192 tokens per request: 73.34x
    (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) INFO 2026-10-13 10:36:59,851 default_LLMServer:google--gemma-4-26B-A4B-it b2grwv9c -- Started vLLM engine.
    INFO 2026-10-13 10:37:05,997 serve 1456 -- Application 'default' is ready at http://127.0.0.1:8000/.
    2026-10-13 10:37:11,447 SUCC cli.py:66 -- ------------------------------------------
    2026-10-13 10:37:11,447 SUCC cli.py:67 -- Job 'raysubmit_94DTwVzR2pvGQVwR' succeeded
    2026-10-13 10:37:11,447 SUCC cli.py:68 -- ------------------------------------------
    

    Ironwood(TPU7x)

    (ServeController pid=2042) INFO 2026-10-13 10:24:38,092 controller 2042 -- Deploying new version of Deployment(name='LLMServer:google--gemma-4-26B-A4B-it', app='default') (initial target replicas: 1).
    (ServeController pid=2042) INFO 2026-10-13 10:24:38,093 controller 2042 -- Deploying new version of Deployment(name='OpenAiIngress', app='default') (initial target replicas: 1).
    (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) INFO 10-13 10:24:49 [__init__.py:318] TPU info: node_name=tpu7x-16-0 | tpu_type=tpu7x-16 | worker_id=0 | num_chips=8 | num_cores_per_chip=2
    (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) INFO 10-13 10:24:49 [__init__.py:238] Platform plugin torchtpu is activated
    (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) INFO 10-13 10:25:23 [core.py:112] Initializing a V1 LLM engine (v0.22.1) with config: model='google/gemma-4-26B-A4B-it', speculative_config=None, tokenizer='google/gemma-4-26B-A4B-it', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.bfloat16, max_seq_len=8192, download_dir=None, load_format=auto, tensor_parallel_size=16, pipeline_parallel_size=1, data_parallel_size=1, decode_context_parallel_size=1, dcp_comm_backend=ag_rs, disable_custom_all_reduce=True, quantization=None, quantization_config=None, enforce_eager=False, enable_return_routed_experts=False, kv_cache_dtype=auto, device_config=tpu, ...
    (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) (RayWorkerWrapper pid=602, ip=172.16.10.6) INFO 10-13 10:25:44 [expert_map_manager.py:245] [EP Rank 8/16] Expert parallelism is enabled. Expert placement strategy: linear. Local/global number of experts: 8/128. Experts local to global index map: 0->64, 1->65, 2->66, 3->67, 4->68, 5->69, 6->70, 7->71.
    (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) (RayWorkerWrapper pid=762) INFO 10-13 10:29:06 [default_loader.py:397] Loading weights took 11.69 seconds
    (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) (RayWorkerWrapper pid=604, ip=172.16.10.6) INFO 10-13 10:34:51 [tpu_worker.py:376] Memory statistics | total_hbm_limit_gb=94.75GiB | total_hbm_limit_cap_gb=87.17GiB | total_hbm_used_gb=3.74GiB | kv_cache_headroom_gb=5.0GiB | kv_connector_hbm_reserve_gb=0.0GiB | total_hbm_avail_gb=78.43GiB
    (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) INFO 10-13 10:34:58 [kv_cache_utils.py:1733] GPU KV cache size: 2,349,568 tokens
    (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) INFO 10-13 10:34:58 [kv_cache_utils.py:1734] Maximum concurrency for 8,192 tokens per request: 286.81x
    (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) INFO 2026-10-13 10:40:13,586 default_LLMServer:google--gemma-4-26B-A4B-it a15k5cv3 -- Started vLLM engine.
    INFO 2026-10-13 10:40:18,415 serve 1953 -- Application 'default' is ready at http://127.0.0.1:8000/.
    2026-10-13 10:40:24,372 SUCC cli.py:66 -- ------------------------------------------
    2026-10-13 10:40:24,372 SUCC cli.py:67 -- Job 'raysubmit_mzWQDQGDM9QSgcMN' succeeded
    2026-10-13 10:40:24,372 SUCC cli.py:68 -- ------------------------------------------
    

サービングされたモデルにクエリを実行する

Ray Serve アプリケーションが RUNNING 状態になったら、登録済みのモデル エンドポイントを確認し、ポート 8000 を介して OpenAI 互換の /v1/chat/completions リクエストを送信します。必要に応じて、スループットとレイテンシのベンチマークを行います。

  1. モデルが /v1/models エンドポイントに登録されていることを確認します。

    curl http://localhost:8000/v1/models
    

    出力は次のようになります。

    {"data":[{"id":"google/gemma-4-26B-A4B-it","object":"model","owned_by":"organization-owner","permission":[],"metadata":{"model_id":"google/gemma-4-26B-A4B-it","input_modality":"text","max_request_context_length":8192}}],"object":"list"}
    
  2. チャット完了プロンプトを /v1/chat/completions エンドポイントに送信します。

    curl -X POST http://localhost:8000/v1/chat/completions \
        -H "Content-Type: application/json" \
        -d '{
            "model": "google/gemma-4-26B-A4B-it",
            "messages": [
                {"role": "system", "content": "You are a helpful AI assistant."},
                {"role": "user", "content": "Explain in one sentence how tensor parallelism and expert parallelism accelerate Mixture-of-Experts LLM inference on TPUs."}
            ],
            "max_tokens": 128,
            "temperature": 0.0
        }'
    

    Ironwood(TPU7x)の出力は次のようになります。

    {"id":"chatcmpl-ec30e61c-011d-49cd-989a-9ae9f59448a8","object":"chat.completion","created":1791887860,"model":"google/gemma-4-26B-A4B-it","choices":[{"index":0,"message":{"role":"assistant","content":"Tensor parallelism accelerates inference by partitioning large weight matrices across TPU cores to reduce per-device memory and computation, while expert parallelism distributes the specialized MoE layers across different devices to enable massive model scaling and parallelize the processing of different tokens through different experts.","refusal":null,"annotations":null,"audio":null,"function_call":null,"tool_calls":[],"reasoning":null},"logprobs":null,"finish_reason":"stop","stop_reason":106,"token_ids":null,"routed_experts":null}],"service_tier":null,"system_fingerprint":"vllm-0.22.1-tp16-ep-aaf76e1a","usage":{"prompt_tokens":49,"total_tokens":100,"completion_tokens":51,"prompt_tokens_details":null},"prompt_logprobs":null,"prompt_token_ids":null,"prompt_text":null,"kv_transfer_params":null}
    
  3. 省略可: vllm bench serve コマンドを使用して、サービングのスループットとレイテンシをベンチマークします。次のコマンドは、Ray TPU ワーカー Pod から vllm bench serve HTTP 負荷生成クライアントを実行し、ポート 8000 で Ray Serve HTTP プロキシに 64 個の同時リクエストを送信します。

    WORKER_POD=$(kubectl get pods -l ray.io/node-type=worker -o jsonpath='{.items[0].metadata.name}')
    kubectl exec -it ${WORKER_POD} -- vllm bench serve \
        --backend openai-chat \
        --endpoint /v1/chat/completions \
        --model google/gemma-4-26B-A4B-it \
        --dataset-name random \
        --random-input-len 1024 \
        --random-output-len 512 \
        --max-concurrency 32 \
        --num-prompts 64 \
        --num-warmups 4 \
        --temperature 0.0 \
        --ignore-eos \
        --host ray-serve-torchtpu-head-svc \
        --port 8000
    

    出力は次のようになります。

    Trillium(v6e)

    ============ Serving Benchmark Result ============
    Successful requests:                     64
    Failed requests:                         0
    Maximum request concurrency:             32
    Benchmark duration (s):                  8.41
    Total input tokens:                      66392
    Total generated tokens:                  32768
    Request throughput (req/s):              7.61
    Output token throughput (tok/s):         3894.26
    Peak output token throughput (tok/s):    4409.00
    Peak concurrent requests:                64.00
    Total token throughput (tok/s):          11784.50
    ---------------Time to First Token----------------
    Mean TTFT (ms):                          414.57
    Median TTFT (ms):                        348.94
    P99 TTFT (ms):                           640.26
    -----Time per Output Token (excl. 1st token)------
    Mean TPOT (ms):                          7.37
    Median TPOT (ms):                        7.29
    P99 TPOT (ms):                           8.24
    ---------------Inter-token Latency----------------
    Mean ITL (ms):                           7.37
    Median ITL (ms):                         0.00
    P99 ITL (ms):                            25.94
    ==================================================
    

    Ironwood(TPU7x)

    ============ Serving Benchmark Result ============
    Successful requests:                     64
    Failed requests:                         0
    Maximum request concurrency:             32
    Benchmark duration (s):                  9.65
    Total input tokens:                      66392
    Total generated tokens:                  32768
    Request throughput (req/s):              6.64
    Output token throughput (tok/s):         3397.33
    Peak output token throughput (tok/s):    4574.00
    Peak concurrent requests:                64.00
    Total token throughput (tok/s):          10280.73
    ---------------Time to First Token----------------
    Mean TTFT (ms):                          654.11
    Median TTFT (ms):                        522.30
    P99 TTFT (ms):                           1288.23
    -----Time per Output Token (excl. 1st token)------
    Mean TPOT (ms):                          8.13
    Median TPOT (ms):                        8.08
    P99 TPOT (ms):                           9.69
    ---------------Inter-token Latency----------------
    Mean ITL (ms):                           8.12
    Median ITL (ms):                         9.31
    P99 ITL (ms):                            26.85
    ==================================================
    

RayService カスタム リソースを使用して宣言的にデプロイする

本番環境のデプロイでは、ジョブを手動で送信する代わりに、マルチホスト Ray クラスタと Ray Serve LLM アプリケーションの両方を単一の宣言型 RayService カスタム リソースとして管理できます。Ray Serve LLM ビルダーは Serve 構成でモデル構成を直接受け取るため、個別の Python スクリプトは必要ありません。KubeRay は、クラスタをプロビジョニングし、Serve アプリケーションをデプロイして、レプリカの健全性をモニタリングします。アプリケーションがトラフィックを処理する準備が整うと、ポート 8000 に Kubernetes ClusterIP Service を作成します。

  1. 前のセクションの RayCluster カスタム リソースがまだ実行されている場合は、ポート転送を停止し、RayCluster カスタム リソースを削除して TPU スライスを解放します。

    pkill -f "kubectl port-forward" || true
    kubectl delete raycluster ray-serve-torchtpu --ignore-not-found
    
  2. サンプル リポジトリで、TPU バージョンの RayService マニフェストを確認します。

    Trillium(v6e)

    apiVersion: ray.io/v1
    kind: RayService
    metadata:
      name: ray-serve-torchtpu
    spec:
      serveConfigV2: |
        applications:
        - name: default
          route_prefix: /
          import_path: ray.serve.llm:build_openai_app
          args:
            llm_configs:
            - model_loading_config:
                model_id: google/gemma-4-26B-A4B-it
              accelerator_type: TPU-V6E
              accelerator_config:
                kind: tpu
                topology: 4x4
              placement_group_config:
                bundle_per_worker:
                  TPU: 1
              engine_kwargs:
                tensor_parallel_size: 16
                enable_expert_parallel: true
                enable_prefix_caching: true
                language_model_only: true
                disable_hybrid_kv_cache_manager: true
                max_model_len: 8192
                max_num_batched_tokens: 4096
              experimental_configs:
                stream_batching_interval_ms: 10
              runtime_env:
                env_vars:
                  TPU_MULTIHOST_BACKEND: "ray"
                  # Execute Gemma 4 routed expert layers on TPU TensorCores instead of SparseCore.
                  USE_MOE_SPARSE_CORE: "0"
                  VLLM_RAY_EXTRA_ENV_VARS_TO_COPY: "USE_MOE_SPARSE_CORE"
      rayClusterConfig:
        headGroupSpec:
          rayStartParams:
            num-cpus: "0"
          template:
            spec:
              containers:
              - name: ray-head
                image: ${IMAGE_URI}
                env:
                - name: HF_TOKEN
                  valueFrom:
                    secretKeyRef:
                      name: hf-secret
                      key: hf_api_token
                ports:
                - containerPort: 6379
                  name: gcs
                - containerPort: 8265
                  name: dashboard
                - containerPort: 10001
                  name: client
                - containerPort: 8000
                  name: serve
                resources:
                  limits:
                    cpu: "4"
                    memory: "16Gi"
                  requests:
                    cpu: "4"
                    memory: "16Gi"
        workerGroupSpecs:
        - groupName: v6e-16
          replicas: 1
          numOfHosts: 4
          template:
            spec:
              containers:
              - name: ray-worker
                image: ${IMAGE_URI}
                env:
                - name: HF_TOKEN
                  valueFrom:
                    secretKeyRef:
                      name: hf-secret
                      key: hf_api_token
                resources:
                  limits:
                    cpu: "64"
                    memory: "200Gi"
                    google.com/tpu: "4"
                  requests:
                    cpu: "64"
                    memory: "200Gi"
                    google.com/tpu: "4"
                volumeMounts:
                - mountPath: /dev/shm
                  name: dshm
              volumes:
              - name: dshm
                emptyDir:
                  medium: Memory
                  sizeLimit: 100Gi
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                cloud.google.com/gke-tpu-topology: 4x4

    Ironwood(TPU7x)

    apiVersion: ray.io/v1
    kind: RayService
    metadata:
      name: ray-serve-torchtpu
    spec:
      serveConfigV2: |
        applications:
        - name: default
          route_prefix: /
          import_path: ray.serve.llm:build_openai_app
          args:
            llm_configs:
            - model_loading_config:
                model_id: google/gemma-4-26B-A4B-it
              accelerator_type: TPU-V7X
              accelerator_config:
                kind: tpu
                topology: 2x2x2
              placement_group_config:
                bundle_per_worker:
                  TPU: 1
              engine_kwargs:
                tensor_parallel_size: 16
                enable_expert_parallel: true
                enable_prefix_caching: true
                language_model_only: true
                disable_hybrid_kv_cache_manager: true
                max_model_len: 8192
                max_num_batched_tokens: 4096
              experimental_configs:
                stream_batching_interval_ms: 10
              runtime_env:
                env_vars:
                  TPU_MULTIHOST_BACKEND: "ray"
                  # Execute Gemma 4 routed expert layers on TPU TensorCores instead of SparseCore.
                  USE_MOE_SPARSE_CORE: "0"
                  VLLM_RAY_EXTRA_ENV_VARS_TO_COPY: "USE_MOE_SPARSE_CORE"
      rayClusterConfig:
        headGroupSpec:
          rayStartParams:
            num-cpus: "0"
          template:
            spec:
              containers:
              - name: ray-head
                image: ${IMAGE_URI}
                env:
                - name: HF_TOKEN
                  valueFrom:
                    secretKeyRef:
                      name: hf-secret
                      key: hf_api_token
                - name: RAY_TPU_RESOURCE_PER_CHIP
                  value: "2"
                ports:
                - containerPort: 6379
                  name: gcs
                - containerPort: 8265
                  name: dashboard
                - containerPort: 10001
                  name: client
                - containerPort: 8000
                  name: serve
                resources:
                  limits:
                    cpu: "4"
                    memory: "16Gi"
                  requests:
                    cpu: "4"
                    memory: "16Gi"
        workerGroupSpecs:
        - groupName: tpu7x-16
          replicas: 1
          numOfHosts: 2
          resources:
            # Advertises 8 logical TPU cores (2 chiplets per physical chip) per host.
            TPU: "8"
          template:
            spec:
              containers:
              - name: ray-worker
                image: ${IMAGE_URI}
                env:
                - name: HF_TOKEN
                  valueFrom:
                    secretKeyRef:
                      name: hf-secret
                      key: hf_api_token
                - name: RAY_TPU_RESOURCE_PER_CHIP
                  value: "2"
                resources:
                  limits:
                    cpu: "64"
                    memory: "250Gi"
                    google.com/tpu: "4"
                  requests:
                    cpu: "64"
                    memory: "250Gi"
                    google.com/tpu: "4"
                volumeMounts:
                - mountPath: /dev/shm
                  name: dshm
              volumes:
              - name: dshm
                emptyDir:
                  medium: Memory
                  sizeLimit: 128Gi
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu7x
                cloud.google.com/gke-tpu-topology: 2x2x2
  3. RayService マニフェストを適用します。

    Trillium(v6e)

    envsubst < ray-service.tpu-v6e-multihost.yaml | kubectl apply -f -
    

    Ironwood(TPU7x)

    envsubst < ray-service.tpu-v7x-multihost.yaml | kubectl apply -f -
    
  4. RayService カスタム リソースがモデルの重みの読み込みと TPU スライス全体での XLA グラフのコンパイルを完了し、Running 状態になるまで待ちます。

    kubectl get rayservice ray-serve-torchtpu
    kubectl get svc ray-serve-torchtpu-serve-svc
    

    出力は次のようになります。

    NAME                 SERVICE STATUS   NUM SERVE ENDPOINTS
    ray-serve-torchtpu   Running          3
    NAME                           TYPE        CLUSTER-IP       EXTERNAL-IP   PORT(S)    AGE
    ray-serve-torchtpu-serve-svc   ClusterIP   34.118.239.213   <none>        8000/TCP   29s
    
  5. ローカルポート 8000 を ray-serve-torchtpu-serve-svc Service に転送し、サービングされたモデルをクエリするの curl コマンドを使用してモデルをクエリします。

    kubectl port-forward svc/ray-serve-torchtpu-serve-svc 8000:8000 >/dev/null 2>&1 &
    

モデルのパフォーマンスをモニタリングする

Cloud Monitoring で TPU TensorCore 使用率、HBM 使用率、Ray クラスタ指標をモニタリングするには、Ray on GKE クラスタのログと指標を収集して表示すると Ray ダッシュボードで TPU 指標を表示するをご覧ください。

Ray ダッシュボードにアクセスする

組み込みの Ray ダッシュボードで Ray Serve デプロイ、レプリカ アクター、ノードラベル、クエリ レイテンシのステータスを調べるには、次の操作を行います。

  1. Ray ヘッドサービスへのポート転送がまだ実行されていない場合は、バックグラウンドで開始します。

    kubectl port-forward svc/ray-serve-torchtpu-head-svc 8265:8265 8000:8000 >/dev/null 2>&1 &
    
  2. Cloud Shell で、[ウェブでプレビュー] アイコン。 [ウェブでプレビュー] をクリックし、[ポート 8265 でプレビュー] を選択するか、ローカル ブラウザで http://localhost:8265 を開きます。

  3. [Serve]、[Cluster]、[Metrics] タブをクリックして、vLLM デプロイ、レプリカの健全性、自動検出された ray.io/tpu-* ノードラベル、リクエスト レイテンシ指標、[TPU Utilization by Node] グラフを表示します。

クリーンアップ

このチュートリアルで使用したリソースについて、 Google Cloud アカウントに課金されないようにするには、RayService または RayCluster カスタム リソース、GKE クラスタ、Artifact Registry リポジトリを削除します。

  1. バックグラウンドのポート転送プロセスを停止し、Ray カスタム リソースと Kubernetes Secret を削除します。

    pkill -f "kubectl port-forward" || true
    kubectl delete rayservice ray-serve-torchtpu --ignore-not-found
    kubectl delete raycluster ray-serve-torchtpu --ignore-not-found
    kubectl delete secret hf-secret --ignore-not-found
    
  2. GKE クラスタを削除します。

    Autopilot

    gcloud container clusters delete ${CLUSTER_NAME} \
        --region=${REGION}
    

    標準

    gcloud container clusters delete ${CLUSTER_NAME} \
        --zone=${ZONE}
    
  3. Artifact Registry リポジトリを削除します。

    gcloud artifacts repositories delete ${ARTIFACT_REGISTRY} \
        --location=${REGION}
    

次のステップ