TPU v6e で vLLM TorchTPU を使用して GKE で Qwen3-30B-A3B をサービングする

このチュートリアルでは、単一ホストの TPU v6e（Trillium）アクセラレータと、TorchTPU（torchtpu）を備えた vLLM TPU サービング フレームワークを使用して、Google Kubernetes Engine（GKE）に Mixture of Experts モデル Qwen3-30B-A3B（Qwen/Qwen3-30B-A3B-FP8）をデプロイして提供する方法について説明します。ワークロードは、データ並列処理とエキスパート並列処理を有効にした 2x2 トポロジの 4 つの TPU v6e チップを含む単一の ct6e-standard-4t ノードで実行します。

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

目標

  • GKE クラスタと単一ホスト TPU v6e ノードプールを設定します。
  • GKE に単一ホストの vLLM TorchTPU サーバーをデプロイします。
  • vLLM TorchTPU サーバーに推論リクエストを送信します。
  • サービング ベンチマーク ワークロードを実行します。
  • クリーンアップする。

費用

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

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

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

始める前に

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

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

    Kubernetes Engine API を有効にする

  • Google Cloud プロジェクトで Artifact Registry API を有効にします。

    Artifact Registry API を有効にする

  • Google Cloud CLI をインストールして初期化し、kubectl をインストールします。Cloud Shell を使用する場合、どちらのツールもプリインストールされています。

  • TPU 容量が配置されているリージョンに、プロジェクトの SSD 割り当てと外部 IP アドレス割り当てが十分にあることを確認します。

    1. Google Cloud コンソールで、[割り当てとシステム上限] ページに移動します。

      [割り当てとシステム上限] に移動

    2. [Persistent Disk SSD (GB)] でフィルタし、モデル キャッシュ ボリュームに 500 GB 以上の割り当てがあることを確認します。

    3. [使用中のリージョン外部 IPv4 アドレス] でフィルタし、ターゲット リージョンに十分な IP アドレスの割り当てがあることを確認します。

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

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

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

必要なロール

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

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

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

環境変数を設定する

認証を設定します。

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="tpu7x-standard-4t"
export NODE_POOL_NAME="np-0"
export RESERVATION_NAME="YOUR_RESERVATION_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export NETWORK="YOUR_NETWORK_NAME"
export SUBNETWORK="YOUR_SUBNETWORK_NAME"
export VLLM_IMAGE="YOUR_VLLM_TPU_IMAGE"

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 クラスタの名前（tpu-v6e-tutorial-cluster など）。
  • YOUR_RESERVATION_NAME: 予約を使用している場合は、TPU v6e 予約の名前。
  • YOUR_HF_TOKEN: Hugging Face アクセス トークン。
  • YOUR_NETWORK_NAME:Google Cloud ネットワークの名前。カスタム ネットワークを作成した場合は、そのネットワークを指定します。それ以外の場合は、default を指定します。
  • YOUR_SUBNETWORK_NAME:Google Cloud サブネットワークの名前。カスタム サブネットワークを作成した場合は、そのサブネットワークを指定します。それ以外の場合は、default を指定します。
  • YOUR_VLLM_TPU_IMAGE: 使用する vLLM TorchTPU イメージ。公開イメージについては、vLLM TorchTPU リポジトリをご覧ください。

GKE クラスタを作成する

Standard モードで GKE クラスタを作成し、単一ホストの TPU v6e ノードプールを接続するには、次の操作を行います。

  1. リージョン GKE クラスタを作成します。

    gcloud container clusters create "${CLUSTER_NAME}" \
    --project="${PROJECT_ID}" \
    --location="${REGION}" \
    --node-locations="${ZONE}" \
    --network="${NETWORK}" \
    --subnetwork="${SUBNETWORK}"

    GKE クラスタのデフォルトのマシンタイプ（e2-medium）には、ワークロード コントローラに十分な CPU 容量がないため、このコマンドでは e2-standard-4 を指定します。

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

    NAME                      LOCATION     MASTER_VERSION      MASTER_IP      MACHINE_TYPE   NODE_VERSION        NUM_NODES  STATUS
tpu-v6e-tutorial-cluster  us-central1  1.35.7-gke.1027000  34.63.222.246  e2-standard-4  1.35.7-gke.1027000  3          RUNNING

  2. 予約を使用して、1 つの ct6e-standard-4t ノード（2x2 トポロジの 4 つの TPU v6e チップ）を持つ単一ホスト TPU v6e ノードプールを作成します。

    gcloud container node-pools create "${NODE_POOL_NAME}" \
    --cluster="${CLUSTER_NAME}" \
    --machine-type="${MACHINE_TYPE}" \
    --location="${REGION}" \
    --node-locations="${ZONE}" \
    --project="${PROJECT_ID}" \
    --num-nodes=1 \
    --reservation-affinity=specific \
    --reservation="${RESERVATION_NAME}"

    特定の予約を使用する代わりに、次のフラグを渡して、Spot 容量または Flex Start VM を使用してノードプールを作成することもできます。

    • Spot VM: --reservation-affinity フラグと --reservation フラグを --spot に置き換えて、Spot 容量を使用します。Spot 容量を使用するインスタンスは、ノードプールの作成後に作成されるまでに時間がかかることがあり、ワークロードの実行中にプリエンプトされる可能性があります。Pod の作成に時間がかかる場合や、ワークロードが実行中に停止する場合は、これが原因である可能性があります。
    • Flex Start VM: Flex Start 消費オプションを使用してノードプールをプロビジョニングするには、--flex-start フラグを渡します。詳細については、TPU と Flex Start VM で小規模なバッチ ワークロードを実行するをご覧ください。

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

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

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

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

認証情報の構成

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 -

vLLM TorchTPU サーバーをデプロイする

ワークロード マニフェスト（qwen3_5_server.yaml）は、次のリソースを定義します。

  • ダウンロードしたモデルの重み、XLA コンパイル キャッシュ、PyTorch コンパイル キャッシュを保存するための、Google Cloud Hyperdisk Balanced ストレージを基盤とする StorageClass（hyperdisk-balanced-qwen3）と 500 GiB の PersistentVolumeClaim（hd-claim-qwen3）。

  • 2x2 トポロジ（google.com/tpu: '4'）を持つ単一の TPU v6e ノードにスケジュールされた Deployment（vllm-qwen3）。コンテナは、次のフラグを使用して Qwen/Qwen3-30B-A3B-FP8 の vllm serve を実行します。

    • --data-parallel-size=4
    • --tensor-parallel-size=1
    • --enable-expert-parallel
    • --kv-cache-dtype=fp8
    • --max-model-len=9236
    • --max-num-seqs=64。

  • ポート 8000 を公開する ClusterIP Service（vllm-qwen3-service）。

ワークロードをデプロイしてモニタリングするには、次の手順を完了します。

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

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

    または、次の内容の qwen3_5_server.yaml という名前のファイルを作成することもできます。

    apiVersion: storage.k8s.io/v1
kind: StorageClass
metadata:
  name: hyperdisk-balanced-qwen3-5
provisioner: pd.csi.storage.gke.io
parameters:
  type: hyperdisk-balanced
reclaimPolicy: Delete
volumeBindingMode: WaitForFirstConsumer
allowVolumeExpansion: true
---
apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: hd-claim-qwen3-5
  namespace: default
spec:
  storageClassName: hyperdisk-balanced-qwen3-5
  accessModes:
    - ReadWriteOnce
  resources:
    requests:
      storage: 1500Gi
---
apiVersion: apps/v1
kind: Deployment
metadata:
  name: vllm-qwen3-5
  namespace: default
spec:
  replicas: 1
  selector:
    matchLabels:
      app: vllm-qwen3-5
  template:
    metadata:
      labels:
        app: vllm-qwen3-5
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu7x
        cloud.google.com/gke-tpu-topology: 2x2x1
      tolerations:
      - key: "cloud.google.com/gke-spot"
        operator: "Equal"
        value: "true"
        effect: "NoSchedule"
      containers:
      - name: vllm-tpu
        image: $VLLM_IMAGE
        command: ["vllm", "serve"]
        args:
        - --host=0.0.0.0
        - --port=8000
        - --model=Qwen/Qwen3.5-397B-A17B-FP8
        - --tokenizer=Qwen/Qwen3.5-397B-A17B-FP8
        - --max-model-len=9236
        - --max-num-batched-tokens=1024
        - --max-num-seqs=64
        - --no-enable-prefix-caching
        - --gpu-memory-utilization=0.9
        - --data-parallel-size=8
        - --tensor-parallel-size=1
        - --async-scheduling
        - --language-model-only
        - --enable-auto-tool-choice
        - --tool-call-parser=qwen3_coder
        - --reasoning-parser=qwen3
        - '--limit-mm-per-prompt={"image": 0, "video": 0}'
        - --kv-cache-dtype=fp8
        - --enable-expert-parallel
        - --download-dir=/data
        - --block-size=128
        env:
        - name: MODEL_IMPL_TYPE
          value: "vllm"
        - name: DP_SCHED_ENABLED
          value: "1"
        - name: ONEHOT_MOE_PERMUTE_THRESHOLD
          value: "32768"
        - name: VLLM_ENGINE_READY_TIMEOUT_S
          value: "9600"
        - name: VLLM_ENGINE_ITERATION_TIMEOUT_S
          value: "1800"
        - name: USE_MOE_SPARSE_CORE
          value: "1"
        - name: TPU_MOE_SKIP_PADDED_TOKENS
          value: "1"
        - name: VLLM_CACHE_ROOT
          value: "/data/vllm_cache"
        - name: VLLM_TORCH_COMPILE_CACHE_DIR
          value: "/data/torch_compile_cache"
        - name: XLA_CACHE_DIR
          value: "/data/xla_cache"
        - name: HF_HOME
          value: /data
        - name: HF_TOKEN
          valueFrom:
            secretKeyRef:
              name: hf-secret
              key: hf_api_token
        ports:
        - containerPort: 8000
        resources:
          limits:
            google.com/tpu: '4'
          requests:
            google.com/tpu: '4'
        readinessProbe:
          tcpSocket:
            port: 8000
          initialDelaySeconds: 15
          periodSeconds: 10
        volumeMounts:
        - mountPath: "/data"
          name: data-volume
        - mountPath: /dev/shm
          name: dshm
      volumes:
      - emptyDir:
          medium: Memory
        name: dshm
      - name: data-volume
        persistentVolumeClaim:
          claimName: hd-claim-qwen3-5
---
apiVersion: v1
kind: Service
metadata:
  name: vllm-qwen3-5-service
  namespace: default
spec:
  selector:
    app: vllm-qwen3-5
  type: ClusterIP
  ports:
  - name: http
    protocol: TCP
    port: 8000
    targetPort: 8000

  2. vLLM TorchTPU サーバー マニフェストをクラスタに適用します。

    envsubst '$VLLM_IMAGE' < "qwen3_5_server.yaml" | kubectl apply -f -

  3. アプリケーションの起動が完了するまで、サーバーのデプロイログをモニタリングします。

    kubectl wait --for=jsonpath='{.status.phase}'=Running pod -l app=vllm-qwen3-5 --timeout=900s
kubectl logs deployment/vllm-qwen3-5 -c vllm-tpu -f

    Pod がモデルの重みをダウンロードしてモデルをコンパイルする間、サービスの初期起動に約 15 ～ 30 分かかることがあります。

    Spot や Flex Start などのプリエンプティブ容量を使用している場合、ノードプールの作成後に TPU スライスノードが作成されるまでに時間がかかることがあります。また、ワークロードの実行中にプリエンプトされる可能性があります。Pod の起動に想定よりも時間がかかっている場合や、ワークロードが途中で停止する場合は、これが原因である可能性があります。詳細については、Spot 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.

  4. 2 つ目のターミナル セッションを開きます。

  5. 2 番目のターミナル セッションで、ローカルポート 8000 を vllm-qwen3-service に転送します。

    kubectl port-forward service/vllm-qwen3-5-service 8000:8000

    このターミナル セッションでポート転送を実行したままにします。

推論リクエストを送信する

最初のターミナル セッションに戻り、/v1/chat/completions エンドポイントにチャット完了リクエストを送信します。

curl http://localhost:8000/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "Qwen/Qwen3.5-397B-A17B-FP8",
        "messages": [
            {"role": "user", "content": "Hello! Tell me a joke."}
        ],
        "max_tokens": 512,
        "temperature": 0.7
    }' | jq .

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

InferenceX ベンチマーク スイートを使用して、8,192 トークンの入力と 1,024 トークンの出力（8k/1k）ワークロードでデプロイされた Qwen3-30B-A3B モデルをベンチマークするには、次の手順を完了します。

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

    curl -O https://raw.githubusercontent.com/GoogleCloudSamples/AIHypercomputerSamples/refs/heads/main/tpu/inference/qwen3-5-397b/qwen3_5-benchmark-8k1k.yaml

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

    apiVersion: v1
kind: Pod
metadata:
  name: qwen3-5-bench-8k1k
  namespace: default
spec:
  restartPolicy: OnFailure
  terminationGracePeriodSeconds: 60
  containers:
  - name: vllm-bench
    image: $VLLM_IMAGE
    command: ["/bin/bash", "-c"]
    args:
    - |
      while ! curl http://vllm-qwen3-5-service:8000/ping; do sleep 30 && echo 'Waiting for server...'; done
      apt-get update && apt-get install -y git && \
      git clone https://github.com/SemiAnalysisAI/InferenceX.git /ubench/inferencex && \
      cd /ubench/inferencex && \
      git checkout 89ce6098ef2bc4576a735c43f39c7d972b091cfc && \
      python3 /ubench/inferencex/utils/bench_serving/benchmark_serving.py \
        --backend=vllm \
        --percentile-metrics='ttft,tpot,itl,e2el' \
        --host=vllm-qwen3-5-service \
        --port=8000 \
        --model=Qwen/Qwen3.5-397B-A17B-FP8 \
        --tokenizer=Qwen/Qwen3.5-397B-A17B-FP8 \
        --dataset-name=random \
        --random-input-len=8192 \
        --random-output-len=1024 \
        --random-range-ratio=1.0 \
        --num-prompts=2560 \
        --num-warmups=512 \
        --max-concurrency=256 \
        --ignore-eos
    env:
    - name: HUGGING_FACE_HUB_TOKEN
      valueFrom:
        secretKeyRef:
          key: hf_api_token
          name: hf-secret

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

    envsubst '$VLLM_IMAGE' < "qwen3_5-benchmark-8k1k.yaml" | kubectl apply -f -

  3. ベンチマーク ログをモニタリングします。

    kubectl wait --for=jsonpath='{.status.phase}'=Running pod --field-selector=metadata.name=qwen3-5-bench-8k1k --timeout=600s
kubectl logs -f qwen3-5-bench-8k1k

    ベンチマークが完了するまでに約 25 ～ 30 分かかります。

  4. ベンチマークが完了したら、次のような指標を含む出力を探します。

    • Output token throughput
    • Total token throughput
    • TTFT
    • TPOT

クリーンアップ

このチュートリアルで使用したリソースについて、 Google Cloud アカウントに継続的な料金が発生しないようにするには、Kubernetes ワークロード、TPU ノードプール、GKE クラスタを削除します。

  1. vLLM TorchTPU サーバーの Deployment、Service、PersistentVolumeClaim、Hugging Face Secret を削除します。

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

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

    gcloud container node-pools delete "${NODE_POOL_NAME}" \
    --cluster="${CLUSTER_NAME}" \
    --location="${REGION}" \
    --project="${PROJECT_ID}" \
    --quiet || true

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

次のステップ