Google は AI 技術を使用して、コンテンツをご希望の言語に翻訳しています。AI 翻訳には誤りが含まれる場合があります。

TPU v6e アクセラレータでマルチホスト vLLM TorchTPU と PD 分離を使用して GKE で Qwen3-30B-A3B をサービングする

このチュートリアルでは、Google Kubernetes Engine（GKE）に Mixture of Experts モデルである Qwen3-30B-A3B（Qwen/Qwen3-30B-A3B-FP8）をデプロイして提供する方法について説明します。このデプロイでは、マルチホスト TPU v6e（Trillium）アクセラレータと、TorchTPU（torchtpu）およびプレフィルデコード（PD）分離を備えた vLLM TPU サービングフレームワークを使用します。

このチュートリアルで説明するデプロイ構成では、2 つのマルチホスト ct6e-standard-4t ノードプールをプロビジョニングします。各ノードプールは 2x4 トポロジで配置されます（ノードプールごとに 2 つのノードに 8 つの TPU v6e チップ、合計 16 個の TPU v6e チップ）。Kubernetes JobSet API と Ray を使用して、次のことを行います。

データ並列処理とエキスパート並列処理を有効にして、分離されたマルチホストプリフィルサーバーとデコードサーバーを別々の TPU v6e スライスで実行します。
リクエストを分離プロキシ経由でルーティングする。

より大きな Qwen3.5-397B モデルを使用して TPU7x（Ironwood）アクセラレータで同様のマルチホスト vLLM TorchTPU サービングワークフローを実行するには、TPU7x アクセラレータでマルチホスト vLLM TorchTPU を使用して GKE で Qwen3.5-397B をサービングするをご覧ください。

目標

GKE クラスタとマルチホスト TPU v6e ノードプールを設定します。
モデルの重みを Cloud Storage バケットにステージングします。
JobSet を使用して、マルチホスト vLLM TorchTPU プリフィル、デコード、プロキシサーバーをデプロイします。
サービングベンチマークワークロードを実行し、結果を分析する。
クリーンアップする。

費用

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

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

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

始める前に

チュートリアルの環境を設定します。

Google Cloud プロジェクトで Kubernetes Engine API を有効にします。

Kubernetes Engine API を有効にする
Google Cloud CLI をインストールして初期化し、kubectl をインストールします。Cloud Shell を使用する場合、どちらのツールもプリインストールされています。
TPU 容量が配置されているリージョンに、プロジェクトの SSD 割り当てと外部 IP アドレス割り当てが十分にあることを確認します。
1. Google Cloud コンソールで、[割り当てとシステム上限] ページに移動します。
  
  [割り当てとシステム上限] に移動
2. Persistent Disk SSD（GB）でフィルタし、ターゲットリージョンに十分な SSD 割り当てがあることを確認します。
3. [使用中のリージョン外部 IPv4 アドレス] でフィルタし、ターゲットリージョンに十分な IP アドレスの割り当てがあることを確認します。

Hugging Face アクセストークンを取得する

このチュートリアルを使用するには、Hugging Face アクセストークンが必要です。Hugging Face で無料アカウントに登録できます。アカウントを作成したら、アクセストークンを生成します。

[Welcome to Hugging Face] ページで、アカウントのアバターをクリックし、[アクセストークン] を選択します。
[アクセストークン] ページで、[Create new token] をクリックします。
[Read] トークンタイプを選択し、トークンの名前を入力します。
アクセストークンが表示されます。トークンを安全な場所に保存します。

必要なロール

このチュートリアルを完了するために必要な権限を取得するには、プロジェクトに対する次の IAM ロールを付与するよう管理者に依頼してください。

Kubernetes Engine 管理者（roles/container.admin）
Compute 管理者（roles/compute.admin）
ストレージ管理者（roles/storage.admin）
サービスアカウントユーザー（roles/iam.serviceAccountUser）

ロールの付与については、プロジェクト、フォルダ、組織へのアクセス権の管理をご覧ください。

必要な権限は、カスタムロールや他の事前定義ロールから取得することもできます。

環境変数を設定する

環境変数を設定する手順は次のとおりです。

認証を設定します。

gcloud auth login
gcloud auth application-default login

このチュートリアル全体で使用する環境変数をエクスポートします。
```
export PROJECT_ID="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export MACHINE_TYPE="ct6e-standard-4t"
export TOPOLOGY="2x4"
export RESERVATION_NAME="YOUR_RESERVATION_NAME"
export GCS_BUCKET="YOUR_GCS_BUCKET"
export HF_TOKEN="YOUR_HF_TOKEN"
export VLLM_IMAGE="vllm/vllm-torchtpu:v0.29.0@sha256:3976ccc1a69ffd361e0a7d2831fe317f6b685636014f426321a39ae02f8da0dc"

gcloud config set project "${PROJECT_ID}"
gcloud config set billing/quota_project "${PROJECT_ID}"
```
次のように置き換えます。
- YOUR_PROJECT_ID: 実際のGoogle Cloud プロジェクト ID。
- YOUR_REGION: クラスタをデプロイするリージョン（us-central1 など）。
- YOUR_ZONE: TPU v6e ノードプールをデプロイするゾーン（us-central1-c など）。
- YOUR_CLUSTER_NAME: GKE クラスタの名前。
- YOUR_RESERVATION_NAME: 予約を使用している場合は、TPU v6e 予約の名前。
- YOUR_GCS_BUCKET: モデルの重みのステージングに使用する Cloud Storage バケットの名前。
- YOUR_HF_TOKEN: Hugging Face アクセストークン。
カスタム vLLM TorchTPU コンテナイメージを使用している場合は、VLLM_IMAGE 変数のデフォルトの vllm/vllm-torchtpu イメージ URI をイメージの URI に置き換えます。公開イメージについては、vLLM TorchTPU リポジトリをご覧ください。

GKE クラスタを作成する

Standard モードで GKE クラスタを作成し、プリフィルとデコードの分離用に 2 つのマルチホスト TPU v6e ノードプールを接続するには、次の操作を行います。

Workload Identity Federation for GKE を有効にしてリージョン GKE クラスタを作成します。

gcloud container clusters create "${CLUSTER_NAME}" \
    --project="${PROJECT_ID}" \
    --location="${REGION}" \
    --node-locations="${ZONE}" \
    --machine-type="e2-standard-4" \
    --workload-pool="${PROJECT_ID}.svc.id.goog"

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

NAME     LOCATION       MASTER_VERSION      MASTER_IP      MACHINE_TYPE   NODE_VERSION        NUM_NODES  STATUS   STACK_TYPE
asq-v6e  us-central1    1.35.7-gke.1027000  35.226.172.84  e2-standard-4  1.35.7-gke.1027000  3          RUNNING  IPV4

2 つのマルチホスト TPU v6e ノードプールを作成します。各ノードプールには 2 つの ct6e-standard-4t ノードが含まれます（ノードプールごとに 2x4 トポロジの 8 つの TPU v6e チップ、合計 16 個の TPU v6e チップ）。

for NODE_POOL_NAME in np-0 np-1; do
  gcloud container node-pools create "${NODE_POOL_NAME}" \
      --cluster="${CLUSTER_NAME}" \
      --machine-type="${MACHINE_TYPE}" \
      --tpu-topology=2x4 \
      --location="${REGION}" \
      --node-locations="${ZONE}" \
      --project="${PROJECT_ID}" \
      --num-nodes=2 \
      --spot
done

このチュートリアルでは、Spot 容量（--spot）を使用します。Spot 容量を使用するインスタンスは、ノードプールの作成後に作成されるまでに時間がかかることがあり、ワークロードの実行中にプリエンプトされる可能性があります。次のフラグを渡して、予約または Flex Start VM を使用してノードプールを作成することもできます。

特定の予約: 特定の予約を使用してノードプールを作成します。

for NODE_POOL_NAME in np-0 np-1; do
  gcloud container node-pools create "${NODE_POOL_NAME}" \
      --cluster="${CLUSTER_NAME}" \
      --machine-type="${MACHINE_TYPE}" \
      --placement-policy="${WORKLOAD_POLICY_NAME}" \
      --location="${REGION}" \
      --node-locations="${ZONE}" \
      --project="${PROJECT_ID}" \
      --num-nodes=2 \
      --reservation-affinity=specific \
      --reservation="${RESERVATION_NAME}"
done

TPU 予約の詳細については、Cloud TPU の予約についてをご覧ください。

Flex Start VM: Flex Start 消費オプションを使用してノードプールをプロビジョニングします。

for NODE_POOL_NAME in np-0 np-1; do
  gcloud container node-pools create "${NODE_POOL_NAME}" \
      --cluster="${CLUSTER_NAME}" \
      --machine-type="${MACHINE_TYPE}" \
      --placement-policy="${WORKLOAD_POLICY_NAME}" \
      --location="${REGION}" \
      --node-locations="${ZONE}" \
      --project="${PROJECT_ID}" \
      --reservation-affinity=none \
      --enable-autoscaling \
      --num-nodes=0 \
      --min-nodes=0 \
      --max-nodes=2 \
      --flex-start
done

詳細については、TPU と Flex Start VM で小規模なバッチワークロードを実行するをご覧ください。

各ノードプールの出力は次のようになります。

NAME  MACHINE_TYPE      DISK_SIZE_GB  NODE_VERSION
np-0  ct6e-standard-4t  100           1.35.7-gke.1027000

クラスタと通信するように kubectl 認証情報を構成します。

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

Cloud Storage にモデルの重みをステージングする

ワークロードの実行用に、モデルの重みを Cloud Storage バケットにステージングします。

均一なバケットレベルのアクセスが有効な Cloud Storage バケットを作成します。

gcloud storage buckets create "gs://${GCS_BUCKET}" \
    --project="${PROJECT_ID}" \
    --location="${REGION}" \
    --uniform-bucket-level-access

バケットに対する Identity and Access Management（IAM）権限を、デフォルトの Compute Engine サービスアカウントとデフォルトの Kubernetes サービスアカウントに付与します。

PROJECT_NUMBER=$(gcloud projects describe "${PROJECT_ID}" --format="value(projectNumber)")

gcloud storage buckets add-iam-policy-binding "gs://${GCS_BUCKET}" \
    --member="serviceAccount:${PROJECT_NUMBER}-compute@developer.gserviceaccount.com" \
    --role="roles/storage.objectAdmin"

gcloud storage buckets add-iam-policy-binding "gs://${GCS_BUCKET}" \
    --member="principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/default/sa/default" \
    --role="roles/storage.admin"

Hugging Face アクセストークンを保存する hf-secret という名前の Kubernetes Secret を作成します。

kubectl create secret generic hf-secret \
    --from-literal=hf_api_token="${HF_TOKEN}" \
    --dry-run=client -o yaml | kubectl apply -f -

download_qwen3_30b.yaml マニフェストを取得します。

curl -O https://raw.githubusercontent.com/GoogleCloudSamples/AIHypercomputerSamples/main/tpu/inference/qwen3-5-397b-multihost/download_qwen3_30b.yaml

または、次の内容のファイルを download_qwen3_30b.yaml という名前で作成します。

apiVersion: batch/v1
kind: Job
metadata:
  name: download-qwen3
  namespace: default
spec:
  backoffLimit: 1
  template:
    spec:
      restartPolicy: OnFailure
      # Schedule onto a TPU node (requests 0 TPUs so it can run alongside or before workloads)
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
      tolerations:
      - operator: Exists
      containers:
      - name: model-downloader
        image: gcr.io/google.com/cloudsdktool/google-cloud-cli:slim
        env:
        - name: HF_TOKEN
          valueFrom:
            secretKeyRef:
              name: hf-secret
              key: hf_api_token
        # Enable high-performance concurrent chunked downloads in huggingface_hub
        - name: HF_XET_HIGH_PERFORMANCE
          value: "1"
        # Place Hugging Face cache on the staging volume so the 100GB boot disk never fills up
        - name: HF_HOME
          value: "/local-ssd/hf-cache"
        command: ["/bin/bash", "-c"]
        args:
        - |
          set -euo pipefail

          echo "=== 1. Installing huggingface_hub ==="
          pip3 install --break-system-packages -q -U "huggingface_hub"

          echo "=== 2. Downloading Qwen/Qwen3-30B-A3B-FP8 to Local Staging Volume ==="
          hf download Qwen/Qwen3-30B-A3B-FP8 --local-dir /local-ssd/model/Qwen/Qwen3-30B-A3B-FP8
          rm -rf /local-ssd/model/Qwen/Qwen3-30B-A3B-FP8/.cache

          echo "=== 3. Uploading from Local Staging Volume to GCS ==="
          gcloud config set storage/parallel_composite_upload_compatibility_check False
          gcloud storage cp -r /local-ssd/model/Qwen gs://${GCS_BUCKET}/

          echo "=== 4. Cleaning up Local Staging Volume ==="
          rm -rf /local-ssd/model /local-ssd/hf-cache
        resources:
          requests:
            cpu: "32"
            memory: "64Gi"
        volumeMounts:
        - name: local-ssd
          mountPath: /local-ssd
      volumes:
      - name: local-ssd
        # If the TPU node pool was created with --ephemeral-storage-local-ssd, use `emptyDir: {}`.
        # On standard ct6e-standard-4t node pools without Local SSDs attached,
        # `emptyDir: {}` uses the 100GB boot disk; use `medium: Memory` (RAM disk) instead:
        emptyDir:
          medium: Memory
          sizeLimit: 64Gi

Qwen/Qwen3-30B-A3B-FP8 モデルの重みをステージングボリュームにダウンロードし、Cloud Storage バケットに転送します。
```
envsubst '$GCS_BUCKET' < download_qwen3_30b.yaml | kubectl apply -f -
```
このステップが完了するまでに約 5 ～ 10 分かかります。

Pod が起動するまで待ち、ダウンロード Job のログをモニタリングします。

kubectl wait --for=condition=Ready pod -l job-name=download-qwen3 --timeout=120s
kubectl logs -f job/download-qwen3

ジョブが完了したら、16 個の .safetensors シャードがすべて Cloud Storage バケットにあることを確認します。
```
gcloud storage ls "gs://${GCS_BUCKET}/Qwen/Qwen3-30B-A3B-FP8/*.safetensors" | wc -l
```
完了したステージングジョブを削除します。
```
kubectl delete job download-qwen3
```

マルチホスト vLLM TorchTPU サーバーをデプロイする

ワークロードマニフェスト（v6e_qwen3_server_multihost_slice.yaml）は、3 つの複製されたジョブを含む Kubernetes JobSet リソース（vllm-qwen3-30b-pd-disagg-v6e）を定義します。これらのジョブは、2 つの 2x4 TPU v6e ノードプール間で分離されたプレフィルとデコードのサービングを調整します。

プリフィルサーバー（prefill）: 1 つの 2x4 TPU v6e ノードプール（ホストあたり 4 つのチップを搭載した 2 つの ct6e-standard-4t ノード、合計 8 つの TPU コア）で 2 つのレプリカを実行します。Pod は 2 ノードの Ray クラスタを形成し、Qwen/Qwen3-30B-A3B-FP8 のポート 8000 で vLLM TorchTPU を起動します。サーバーは、次のフラグを使用して kv_producer として構成された TPURaidenConnector を使用します。
- --data-parallel-size=8
- --tensor-parallel-size=1
- --enable-expert-parallel
- --load-format=runai_streamer
デコードサーバー（decode）: 2 番目の 2x4 TPU v6e ノードプール（合計 8 個の TPU コア）に 2 つのレプリカを実行します。JobSet は、cloud.google.com/gke-nodepool ラベルの Pod アフィニティルールとアンチアフィニティルールを使用して、プリフィルとデコードを別々のノードプールに分離します。Pod は 2 ノードの Ray クラスタを形成し、kv_consumer として構成された TPURaidenConnector を使用してポート 7200 で vLLM TorchTPU を起動します。
分離プロキシ（proxy）: ポート 10000 で toy_proxy_server.py スクリプトを実行し、プリフィルサーバー（ポート 8000）とデコードサーバー（ポート 7200）間のリクエストルーティングと KV キャッシュ転送を調整します。

分離された vLLM サーバーをデプロイしてモニタリングする手順は次のとおりです。

Kubernetes JobSet CustomResourceDefinition とコントローラをクラスタにインストールします。

kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/v0.12.0/manifests.yaml
kubectl wait --for=condition=Available --timeout=300s deployment/jobset-controller-manager -n jobset-system

サーバーマニフェストを取得します。

curl -O https://raw.githubusercontent.com/GoogleCloudSamples/AIHypercomputerSamples/main/tpu/inference/qwen3-5-397b-multihost/v6e_qwen3_server_multihost_slice.yaml

または、次の内容で v6e_qwen3_server_multihost_slice.yaml という名前のファイルを作成します。

apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: vllm-qwen3-30b-pd-disagg-v6e
spec:
  failurePolicy:
    restartStrategy: BlockingRecreate
    maxRestarts: 5
  ttlSecondsAfterFinished: 86400
  network:
    enableDNSHostnames: true
  replicatedJobs:
  # ---------------------------------------------------------
  # PREFILL SERVERS (2 Replicas, 1 Node each = 4 chips on 2x4 v6e slice)
  # ---------------------------------------------------------
  - name: prefill
    replicas: 2
    template:
      spec:
        parallelism: 1
        completions: 1
        backoffLimit: 0
        template:
          metadata:
            labels:
              app: vllm-tpu-prefill
          spec:
            dnsPolicy: ClusterFirstWithHostNet
            nodeSelector:
              cloud.google.com/gke-tpu-topology: 2x4
              cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
            tolerations:
            - key: "google.com/tpu"
              operator: "Exists"
              effect: "NoSchedule"
            - key: "cloud.google.com/gke-spot"
              operator: "Equal"
              value: "true"
              effect: "NoSchedule"
            affinity:
              podAntiAffinity:
                requiredDuringSchedulingIgnoredDuringExecution:
                - labelSelector:
                    matchExpressions:
                    - key: app
                      operator: In
                      values:
                      - vllm-tpu-decode
                  topologyKey: "cloud.google.com/gke-nodepool"
              podAffinity:
                requiredDuringSchedulingIgnoredDuringExecution:
                - labelSelector:
                    matchExpressions:
                    - key: app
                      operator: In
                      values:
                      - vllm-tpu-prefill
                  topologyKey: "cloud.google.com/gke-nodepool"
            initContainers:
            - name: tpu-node-setup
              image: busybox
              securityContext:
                privileged: true
              command:
              - /bin/sh
              - -c
              args:
              - |
                sysctl -w vm.max_map_count=8388608
                if [ -f /sys/module/vfio_iommu_type1/parameters/dma_entry_limit ]; then
                  echo 2000000 > /sys/module/vfio_iommu_type1/parameters/dma_entry_limit
                fi
            containers:
            - name: vllm-tpu-prefill
              image: ${VLLM_IMAGE}
              imagePullPolicy: IfNotPresent
              command:
              - /bin/bash
              - -c
              args:
                - |
                  set -e
                  export TPU_WORKER_ID=$JOBSET_JOB_INDEX
                  unset TPU_VISIBLE_CHIPS

                  echo "=== Installing runai-model-streamer[gcs] and pre-caching metadata from GCS on $HOSTNAME ==="
                  pip install -q "runai-model-streamer[gcs]==0.16.1"
                  python3 -c "from vllm.transformers_utils.runai_utils import ObjectStorageModel; uri = 'gs://${GCS_BUCKET}/Qwen/Qwen3-30B-A3B-FP8'; m = ObjectStorageModel(uri); m.pull_files(uri, ignore_pattern=['*.pt', '*.safetensors', '*.bin', '*.tensors', '*.pth']); print('Pulled metadata/tokenizer to', m.dir)"

                  HEAD_ADDR="${JS_NAME}-prefill-0-0.${JS_NAME}"

                  if [ "$JOBSET_JOB_INDEX" = "0" ]; then
                    echo "=== Node 0: Starting Ray Head on $HOSTNAME (port 6379, IP $POD_IP) ==="
                    ray start --head --port=6379 --node-ip-address=$POD_IP --resources='{"TPU": 4}' --block &

                    echo "Waiting for worker node to join Ray cluster (expecting 8.0 TPU total across 2 nodes)..."
                    until ray status 2>/dev/null | grep -E -q "8(\.0)? TPU"; do
                      echo "Waiting for worker to connect to Ray Head..."
                      sleep 2
                    done
                    echo "Worker joined! Ray cluster status:"
                    ray status

                    echo "=== Starting vLLM Prefill Server (Qwen3 30B FP8, DP=8, TP=1, Port 8000) ==="
                    PREFILL_ARGS=(
                      gs://${GCS_BUCKET}/Qwen/Qwen3-30B-A3B-FP8
                      --served-model-name=Qwen/Qwen3-30B-A3B-FP8
                      --api-server-count=1
                      --load-format=runai_streamer
                      --enable-ep-weight-filter
                      --tensor-parallel-size=1
                      --data-parallel-size=8
                      --data-parallel-size-local=4
                      --data-parallel-backend=ray
                      --data-parallel-address=$POD_IP
                      --distributed-executor-backend=ray
                      --enable-expert-parallel
                      --max-model-len=18432
                      --max-num-batched-tokens=8192
                      --max-num-seqs=64
                      --num-gpu-blocks-override=128
                      --host=0.0.0.0
                      --port=8000
                      --async-scheduling
                      --no-enable-prefix-caching
                      --kv-cache-dtype=fp8
                      --attention-backend=CUSTOM
                      --block-size=1024
                      --return-tokens-as-token-ids
                      --hf-overrides '{"linear_num_key_heads":1,"linear_num_value_heads":1}'
                      --compilation-config '{"backend":"vllm_torchtpu.compilation.tpu_compiler.TpuCompilerAdaptor","compile_sizes":[4096,8192],"inductor_compile_config":{"enable_auto_functionalized_v2":false,"size_asserts":false,"alignment_asserts":false,"scalar_asserts":false}}'
                      --kv-transfer-config '{"kv_connector":"TPURaidenConnector","kv_connector_module_path":"vllm_torchtpu.distributed.kv_transfer.tpu_connector","kv_role":"kv_producer","kv_port":14579}'
                    )
                    exec vllm serve "${PREFILL_ARGS[@]}"
                  else
                    echo "=== Node 1: Starting Ray Worker on $HOSTNAME (IP $POD_IP) ==="
                    until timeout 1 bash -c "cat < /dev/null > /dev/tcp/${HEAD_ADDR}/6379" 2>/dev/null; do
                      echo "Waiting for Ray head at ${HEAD_ADDR}:6379..."
                      sleep 2
                    done
                    echo "Ray head is reachable! Starting Ray worker..."
                    sleep 3
                    exec ray start --address="${HEAD_ADDR}:6379" --node-ip-address=$POD_IP --resources='{"TPU": 4}' --block
                  fi

              env:
              - name: JS_NAME
                valueFrom:
                  fieldRef:
                    fieldPath: metadata.labels['jobset.sigs.k8s.io/jobset-name']
              - name: JOBSET_JOB_INDEX
                valueFrom:
                  fieldRef:
                    fieldPath: metadata.labels['jobset.sigs.k8s.io/job-index']
              - name: POD_IP
                valueFrom:
                  fieldRef:
                    fieldPath: status.podIP
              - name: TPU_MULTIHOST_BACKEND
                value: "ray"
              - name: RAY_EXPERIMENTAL_NOSET_TPU_VISIBLE_CHIPS
                value: "1"
              - name: RAY_ADDRESS
                value: "127.0.0.1:6379"
              - name: VLLM_ENGINE_READY_TIMEOUT_S
                value: "14400"
              - name: SKIP_MEGASCALE_PJRT_CLIENT
                value: "true"
              - name: USE_MOE_FUSED_EP_KERNEL
                value: "1"
              - name: RAGGED_GATHER_REDUCE_VERSION
                value: "v3"
              - name: LIBTPU_INIT_ARGS
                value: "--xla_tpu_use_minor_sharding_for_major_trivial_input=true --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=false --xla_tpu_ars_combiner_threshold_in_bytes=0 --xla_tpu_enable_async_collective_merger=false --xla_tpu_enable_diagnostic_coordinator=false --xla_enable_async_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_all_gather=true --xla_tpu_sparse_core_all_gather_offload_min_size_in_bytes=262144"
              # Raiden KV-transfer configuration (Prefill / Producer)
              - name: TPU_KV_RESHARD_TRANSPORT
                value: "raiden"
              - name: TPU_RAIDEN_RESHARD_IMPL
                value: "store"
              - name: TPU_RAIDEN_ADVERTISE_HOST
                valueFrom:
                  fieldRef:
                    fieldPath: status.podIP
              - name: TPU_USE_RAIDEN_KV_CACHE_MANAGER
                value: "1"
              - name: TPU_RAIDEN_QWEN35_ADMISSION
                value: "1"
              - name: TPU_RAIDEN_TRANSFER_PARALLELISM
                value: "1"
              - name: TPU_VLLM_ENABLE_UNIFIED_BLOCK_POOL
                value: "1"
              volumeMounts:
              - name: dshm
                mountPath: /dev/shm
              securityContext:
                privileged: true
                capabilities:
                  add:
                  - IPC_LOCK
              ports:
              - name: http
                containerPort: 8000
              - name: ray-gcs
                containerPort: 6379
              resources:
                limits:
                  google.com/tpu: "4"
                requests:
                  google.com/tpu: "4"
            volumes:
            - name: dshm
              emptyDir:
                medium: Memory

  # ---------------------------------------------------------
  # DECODE SERVERS (2 Replicas, 1 Node each = 4 chips on 2x4 v6e slice)
  # ---------------------------------------------------------
  - name: decode
    replicas: 2
    template:
      spec:
        parallelism: 1
        completions: 1
        backoffLimit: 0
        template:
          metadata:
            labels:
              app: vllm-tpu-decode
          spec:
            dnsPolicy: ClusterFirstWithHostNet
            nodeSelector:
              cloud.google.com/gke-tpu-topology: 2x4
              cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
            tolerations:
            - key: "google.com/tpu"
              operator: "Exists"
              effect: "NoSchedule"
            - key: "cloud.google.com/gke-spot"
              operator: "Equal"
              value: "true"
              effect: "NoSchedule"
            affinity:
              podAntiAffinity:
                requiredDuringSchedulingIgnoredDuringExecution:
                - labelSelector:
                    matchExpressions:
                    - key: app
                      operator: In
                      values:
                      - vllm-tpu-prefill
                  topologyKey: "cloud.google.com/gke-nodepool"
              podAffinity:
                requiredDuringSchedulingIgnoredDuringExecution:
                - labelSelector:
                    matchExpressions:
                    - key: app
                      operator: In
                      values:
                      - vllm-tpu-decode
                  topologyKey: "cloud.google.com/gke-nodepool"
            initContainers:
            - name: tpu-node-setup
              image: busybox
              securityContext:
                privileged: true
              command:
              - /bin/sh
              - -c
              args:
              - |
                sysctl -w vm.max_map_count=8388608
                if [ -f /sys/module/vfio_iommu_type1/parameters/dma_entry_limit ]; then
                  echo 2000000 > /sys/module/vfio_iommu_type1/parameters/dma_entry_limit
                fi
            containers:
            - name: vllm-tpu-decode
              image: ${VLLM_IMAGE}
              imagePullPolicy: IfNotPresent
              command:
              - /bin/bash
              - -c
              args:
                - |
                  set -e
                  export TPU_WORKER_ID=$JOBSET_JOB_INDEX
                  unset TPU_VISIBLE_CHIPS

                  echo "=== Installing runai-model-streamer[gcs] and pre-caching metadata from GCS on $HOSTNAME ==="
                  pip install -q "runai-model-streamer[gcs]==0.16.1"
                  python3 -c "from vllm.transformers_utils.runai_utils import ObjectStorageModel; uri = 'gs://${GCS_BUCKET}/Qwen/Qwen3-30B-A3B-FP8'; m = ObjectStorageModel(uri); m.pull_files(uri, ignore_pattern=['*.pt', '*.safetensors', '*.bin', '*.tensors', '*.pth']); print('Pulled metadata/tokenizer to', m.dir)"

                  HEAD_ADDR="${JS_NAME}-decode-0-0.${JS_NAME}"

                  if [ "$JOBSET_JOB_INDEX" = "0" ]; then
                    echo "=== Node 0: Starting Ray Head on $HOSTNAME (port 6379, IP $POD_IP) ==="
                    ray start --head --port=6379 --node-ip-address=$POD_IP --resources='{"TPU": 4}' --block &

                    echo "Waiting for worker node to join Ray cluster (expecting 8.0 TPU total across 2 nodes)..."
                    until ray status 2>/dev/null | grep -E -q "8(\.0)? TPU"; do
                      echo "Waiting for worker to connect to Ray Head..."
                      sleep 2
                    done
                    echo "Worker joined! Ray cluster status:"
                    ray status

                    echo "=== Starting vLLM Decode Server (Qwen3 30B FP8, DP=8, TP=1, Port 7200) ==="
                    DECODE_ARGS=(
                      gs://${GCS_BUCKET}/Qwen/Qwen3-30B-A3B-FP8
                      --served-model-name=Qwen/Qwen3-30B-A3B-FP8
                      --api-server-count=1
                      --load-format=runai_streamer
                      --enable-ep-weight-filter
                      --tensor-parallel-size=1
                      --data-parallel-size=8
                      --data-parallel-size-local=4
                      --data-parallel-backend=ray
                      --data-parallel-address=$POD_IP
                      --distributed-executor-backend=ray
                      --enable-expert-parallel
                      --max-model-len=18432
                      --max-num-batched-tokens=2048
                      --max-num-seqs=64
                      --num-gpu-blocks-override=32
                      --host=0.0.0.0
                      --port=7200
                      --async-scheduling
                      --no-enable-prefix-caching
                      --kv-cache-dtype=fp8
                      --attention-backend=CUSTOM
                      --block-size=8192
                      --hf-overrides '{"linear_num_key_heads":1,"linear_num_value_heads":1}'
                      --compilation-config '{"backend":"vllm_torchtpu.compilation.tpu_compiler.TpuCompilerAdaptor","compile_sizes":[8,16,32,64,2048],"inductor_compile_config":{"enable_auto_functionalized_v2":false,"size_asserts":false,"alignment_asserts":false,"scalar_asserts":false}}'
                      --kv-transfer-config '{"kv_connector":"TPURaidenConnector","kv_connector_module_path":"vllm_torchtpu.distributed.kv_transfer.tpu_connector","kv_role":"kv_consumer","kv_port":14579}'
                    )
                    exec vllm serve "${DECODE_ARGS[@]}"
                  else
                    echo "=== Node 1: Starting Ray Worker on $HOSTNAME (IP $POD_IP) ==="
                    until timeout 1 bash -c "cat < /dev/null > /dev/tcp/${HEAD_ADDR}/6379" 2>/dev/null; do
                      echo "Waiting for Ray head at ${HEAD_ADDR}:6379..."
                      sleep 2
                    done
                    echo "Ray head is reachable! Starting Ray worker..."
                    sleep 3
                    exec ray start --address="${HEAD_ADDR}:6379" --node-ip-address=$POD_IP --resources='{"TPU": 4}' --block
                  fi

              env:
              - name: JS_NAME
                valueFrom:
                  fieldRef:
                    fieldPath: metadata.labels['jobset.sigs.k8s.io/jobset-name']
              - name: JOBSET_JOB_INDEX
                valueFrom:
                  fieldRef:
                    fieldPath: metadata.labels['jobset.sigs.k8s.io/job-index']
              - name: POD_IP
                valueFrom:
                  fieldRef:
                    fieldPath: status.podIP
              - name: TPU_MULTIHOST_BACKEND
                value: "ray"
              - name: RAY_EXPERIMENTAL_NOSET_TPU_VISIBLE_CHIPS
                value: "1"
              - name: RAY_ADDRESS
                value: "127.0.0.1:6379"
              - name: VLLM_ENGINE_READY_TIMEOUT_S
                value: "14400"
              - name: SKIP_MEGASCALE_PJRT_CLIENT
                value: "true"
              - name: USE_MOE_FUSED_EP_KERNEL
                value: "1"
              - name: RAGGED_GATHER_REDUCE_VERSION
                value: "v3"
              - name: USE_BATCHED_RPA_LONGCTX
                value: "1"
              - name: LIBTPU_INIT_ARGS
                value: "--xla_tpu_use_minor_sharding_for_major_trivial_input=true --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=false --xla_tpu_ars_combiner_threshold_in_bytes=0 --xla_tpu_enable_async_collective_merger=false --xla_tpu_enable_diagnostic_coordinator=false --xla_enable_async_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_all_gather=true --xla_tpu_sparse_core_all_gather_offload_min_size_in_bytes=262144"
              # Raiden KV-transfer configuration (Decode / Consumer)
              - name: TPU_KV_RESHARD_TRANSPORT
                value: "raiden"
              - name: TPU_RAIDEN_RESHARD_IMPL
                value: "store"
              - name: TPU_RAIDEN_ADVERTISE_HOST
                valueFrom:
                  fieldRef:
                    fieldPath: status.podIP
              - name: TPU_USE_RAIDEN_KV_CACHE_MANAGER
                value: "1"
              - name: TPU_RAIDEN_QWEN35_ADMISSION
                value: "1"
              - name: TPU_RAIDEN_TRANSFER_PARALLELISM
                value: "1"
              - name: TPU_VLLM_ENABLE_UNIFIED_BLOCK_POOL
                value: "1"
              volumeMounts:
              - name: dshm
                mountPath: /dev/shm
              securityContext:
                privileged: true
                capabilities:
                  add:
                  - IPC_LOCK
              ports:
              - name: http
                containerPort: 7200
              - name: ray-gcs
                containerPort: 6379
              resources:
                limits:
                  google.com/tpu: "4"
                requests:
                  google.com/tpu: "4"
            volumes:
            - name: dshm
              emptyDir:
                medium: Memory

  # ---------------------------------------------------------
  # DISAGGREGATION PROXY (Coordinates Prefill -> Decode)
  # ---------------------------------------------------------
  - name: proxy
    replicas: 1
    template:
      spec:
        parallelism: 1
        completions: 1
        backoffLimit: 2
        template:
          spec:
            containers:
            - name: vllm-tpu-proxy
              image: ${VLLM_IMAGE}
              imagePullPolicy: IfNotPresent
              command:
              - /bin/bash
              - -c
              args:
              - |
                set -e
                exec python3 /root/torchtpu-vllm/examples/disagg/toy_proxy_server.py \
                  --host 0.0.0.0 \
                  --port 10000 \
                  --prefiller-host "${JS_NAME}-prefill-0-0.${JS_NAME}" \
                  --prefiller-port 8000 \
                  --decoder-host "${JS_NAME}-decode-0-0.${JS_NAME}" \
                  --decoder-port 7200
              env:
              - name: JS_NAME
                valueFrom:
                  fieldRef:
                    fieldPath: metadata.labels['jobset.sigs.k8s.io/jobset-name']
              ports:
              - name: http
                containerPort: 10000
              resources:
                requests: { cpu: "2", memory: "4Gi" }
                limits: { cpu: "8", memory: "16Gi" }

Cloud Storage バケット名とコンテナイメージを v6e_qwen3_server_multihost_slice.yaml ファイルに置き換え、マニフェストを適用します。
```
envsubst '${GCS_BUCKET} ${VLLM_IMAGE}' < "v6e_qwen3_server_multihost_slice.yaml" | kubectl apply -f -
```
この構成では、プリフィルに 1 つの 2x4 TPU v6e スライス（2 つのノード）を使用し、デコードに 1 つの 2x4 TPU v6e スライス（2 つのノード）を使用します。いずれかのプールのノード数をスケーリングするには、replicas フィールド（replicas: NUM_NODES）を更新します。NUM_NODES は、使用するノード数に置き換えます。
アプリケーションの起動が完了するまで、プリフィル Pod とデコード Pod のログをモニタリングします。
```
# Prefill logs
kubectl logs -l jobset.sigs.k8s.io/replicatedjob-name=prefill -f --tail=-1

# Decode logs
kubectl logs -l jobset.sigs.k8s.io/replicatedjob-name=decode -f --tail=-1
```
最初の実行では、モデルの初期読み込みとコンパイルに約 15 ～ 25 分かかります。

Spot 容量や Flex Start VM などのプリエンプティブル容量を使用している場合、ノードプールの作成後に TPU スライスノードが作成されるまでに時間がかかることがあります。また、ワークロードの実行中にプリエンプトされる可能性があります。詳細については、プリエンプティブル VM インスタンスをご覧ください。TPU スライスノードが作成されているかどうかを確認するには、ノードのステータスと条件を確認するをご覧ください。

プリフィルサーバーログとデコードサーバーログの両方で Application startup complete が報告されるまで待ってから、ベンチマークワークロードを実行します。
```
(APIServer pid=1) INFO:     Started server process [1]
(APIServer pid=1) INFO:     Waiting for application startup.
(APIServer pid=1) INFO:     Application startup complete.
```

ベンチマークワークロードを実行する

8,192 トークンの入力と 1,024 トークンの出力（8k/1k）ワークロードを使用して、デプロイされた Qwen3-30B-A3B モデルを分離プロキシに対してベンチマークするには、次の操作を行います。

ベンチマークマニフェストを取得します。

curl -O https://raw.githubusercontent.com/GoogleCloudSamples/AIHypercomputerSamples/main/tpu/inference/qwen3-5-397b-multihost/v6e_qwen3_benchmark_multihost_slice.yaml

または、次の内容で v6e_qwen3_benchmark_multihost_slice.yaml という名前のファイルを作成します。

apiVersion: batch/v1
kind: Job
metadata:
  name: vllm-qwen3-30b-pd-disagg-v6e-tester
spec:
  parallelism: 1
  completions: 1
  backoffLimit: 0
  template:
    spec:
      restartPolicy: Never
      containers:
      - name: vllm-test
        image: ${VLLM_IMAGE}
        imagePullPolicy: IfNotPresent
        resources:
          requests: { cpu: "1", memory: "4Gi" }
          limits: { cpu: "4", memory: "8Gi" }
        command:
        - /bin/bash
        - -c
        args:
        - |
          set -e
          PREFILLER_DNS="${JS_NAME}-prefill-0-0.${JS_NAME}"
          DECODER_DNS="${JS_NAME}-decode-0-0.${JS_NAME}"
          PROXY_DNS="${JS_NAME}-proxy-0-0.${JS_NAME}"
          MODEL="Qwen/Qwen3-30B-A3B-FP8"

          echo 'import torch; torch.utils.rename_privateuse1_backend("tpu")' > /tmp/sitecustomize.py
          export PYTHONPATH="/tmp:${PYTHONPATH:-}"

          wait_for() {
            local name="$1" url="$2" timeout=18000 elapsed=0
            echo "Waiting for $name ($url)..."
            until curl -sf "$url" > /dev/null; do
              [ "$elapsed" -ge "$timeout" ] && { echo "ERROR: $name timeout" >&2; exit 1; }
              sleep 10; elapsed=$((elapsed + 10))
            done
            echo "$name is up after ${elapsed}s!"
          }
          wait_for prefill "http://$PREFILLER_DNS:8000/health"
          wait_for decode  "http://$DECODER_DNS:7200/health"
          wait_for proxy   "http://$PROXY_DNS:10000/healthcheck"

          echo "=== BENCHMARK: 8K/1K c=128 n=1280 ==="
          BENCH_ARGS=(
            --backend openai-chat
            --base-url "http://$PROXY_DNS:10000" --endpoint /v1/chat/completions
            --model "$MODEL" --dataset-name random
            --random-input-len 8192 --random-output-len 1024
            --ignore-eos
            --max-concurrency 128 --num-prompts 1280
          )
          if vllm bench serve "${BENCH_ARGS[@]}"; then
            echo "=== BENCHMARK COMPLETED SUCCESSFULLY ==="
          else
            echo "ERROR: BENCHMARK FAILED"
            exit 1
          fi
        env:
        - name: JS_NAME
          value: "vllm-qwen3-30b-pd-disagg-v6e"

ベンチマーク Job マニフェストを適用します。

envsubst '${VLLM_IMAGE}' < v6e_qwen3_benchmark_multihost_slice.yaml | kubectl apply -f -

ベンチマーク Pod の準備が整うまで待ち、ベンチマーク Job のログをモニタリングします。

kubectl wait --for=condition=Ready pod -l job-name=vllm-qwen3-30b-pd-disagg-v6e-tester --timeout=300s
kubectl logs -f job/vllm-qwen3-30b-pd-disagg-v6e-tester

ベンチマークが完了したら、次のような指標を含む出力を探します。
- Output token throughput
- Total token throughput
- TTFT
- TPOT

クリーンアップ

このチュートリアルで使用したリソースについて Google Cloud アカウントに課金されないようにするには、Kubernetes ワークロード、TPU v6e ノードプール、GKE クラスタ、Cloud Storage バケットを削除します。

Kubernetes ワークロードを削除します。

kubectl delete -f v6e_qwen3_server_multihost_slice.yaml --ignore-not-found || true
kubectl delete secret hf-secret --ignore-not-found || true

TPU v6e ノードプールと GKE クラスタを削除します。

for NODE_POOL_NAME in np-0 np-1; do
  gcloud container node-pools delete "${NODE_POOL_NAME}" \
      --cluster="${CLUSTER_NAME}" \
      --location="${REGION}" \
      --project="${PROJECT_ID}" \
      --quiet || true
done

gcloud container clusters delete "${CLUSTER_NAME}" \
    --location="${REGION}" \
    --project="${PROJECT_ID}" \
    --quiet || true

ステージングされたモデルの重みを保存する Cloud Storage バケットを削除します。モデルの重みを保持する場合は、この手順をスキップするか、任意の場所にコピーします。
```
gcloud storage rm --recursive "gs://${GCS_BUCKET}" --project="${PROJECT_ID}" --quiet || true
```

次のステップ

Cloud TPU での vLLM の詳細を確認する。
TPU v6e（Trillium）の詳細を確認する。
Cloud TPU の詳細を確認する。

TPU v6e アクセラレータでマルチホスト vLLM TorchTPU と PD 分離を使用して GKE で Qwen3-30B-A3B をサービングする コレクションでコンテンツを整理 必要に応じて、コンテンツの保存と分類を行います。

目標

費用