在 GKE 上使用多主機 vLLM TorchTPU，透過 TPU7x 加速器提供 Qwen3.5-397B

本教學課程說明如何在 Google Kubernetes Engine (GKE) 部署及提供 Qwen3.5-397B (Qwen/Qwen3.5-397B-A17B-FP8)，這是一種專家混合模型。這項部署作業會使用多主機 TPU7x (Ironwood) 加速器和 vLLM TPU 服務架構，並搭配 TorchTPU (torchtpu) 和預填解碼 (PD) 分解。

在這個設定中，您會佈建兩個多主機 tpu7x-standard-4t 節點集區，每個集區都採用 2x2x2 拓撲 (每個節點集區有兩個節點，共八個 TPU7x 晶片，總計 16 個 TPU7x 晶片)。您可以使用 Kubernetes JobSet API 和 Ray，在不同的 TPU7x 磁碟區上執行分散式多主機預填和解碼伺服器，並啟用資料平行處理和專家平行處理，然後透過分散式 Proxy 轉送要求。

目標

  • 設定 GKE 叢集和多主機 TPU7x 節點集區。
  • 將模型權重暫存至 Cloud Storage bucket。
  • 使用 JobSet 部署多主機 vLLM TorchTPU 預先填入、解碼和 Proxy 伺服器。
  • 執行基準測試工作負載，並分析結果。
  • 清除所用資源。

費用

在本文件中，您會使用下列Google Cloud計費元件：

您可以使用 Pricing Calculator 根據預測用量估算費用。

完成本文所述工作後，您可以刪除建立的資源，避免繼續計費，詳情請參閱「清除所用資源」一節。

事前準備

  • 您需要 Hugging Face 存取權杖，才能使用本教學課程。您可以前往 Hugging Face 註冊免費帳戶。擁有帳戶後，請產生存取權杖：

    1. 在「Welcome to Hugging Face」頁面，點選帳戶顯示圖片，然後選取「存取權杖」。
    2. 在「存取權杖」頁面，按一下「建立新權杖」。
    3. 選取「Read」權杖類型，並輸入權杖名稱。
    4. 系統會顯示存取權杖。將權杖儲存在安全的地方。

如要取得完成本教學課程所需的權限，請要求管理員在專案中授予您下列 IAM 角色：

如要進一步瞭解如何授予角色，請參閱「管理專案、資料夾和組織的存取權」。

您或許也能透過自訂角色或其他預先定義的角色，取得必要權限。

設定環境變數

  1. 設定驗證機制：

    gcloud auth login
gcloud auth application-default login

  2. 匯出本教學課程中使用的環境變數：

    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 WORKLOAD_POLICY_NAME="${CLUSTER_NAME}-policy"
export NODE_POOL_NAME_0="np-0"
export NODE_POOL_NAME_1="np-1"
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：要部署 TPU7x 節點集區的可用區，例如 us-central1-c。
    • YOUR_CLUSTER_NAME：GKE 叢集名稱。
    • YOUR_RESERVATION_NAME：TPU 預留項目名稱 (如有使用預留項目)。
    • YOUR_GCS_BUCKET：用於暫存模型權重的 Cloud Storage bucket 名稱。
    • YOUR_HF_TOKEN：您的 Hugging Face 存取權杖。
    • YOUR_VLLM_IMAGE：要使用的 vLLM TorchTPU 映像檔。如要查看公開映像檔，請參閱 vLLM TorchTPU 存放區。

建立 GKE 叢集

如要以標準模式建立 GKE 叢集，並連接兩個多主機 TPU7x 節點集區，用於預先填入和解碼拆解，請完成下列步驟：

  1. 建立區域性 GKE 叢集，並啟用 Workload Identity Federation for 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-v7x  us-central1    1.35.7-gke.1027000  35.226.172.84  e2-standard-4  1.35.7-gke.1027000  3          RUNNING  IPV4

  2. 建立高處理量工作負載配置政策，並採用 2x2x2 加速器拓撲，這是多主機 TPU7x 節點集區的必要條件：

    gcloud compute resource-policies create workload-policy "${WORKLOAD_POLICY_NAME}" \
    --type=HIGH_THROUGHPUT \
    --accelerator-topology=2x2x2 \
    --project="${PROJECT_ID}" \
    --region="${REGION}"

  3. 建立兩個多主機 TPU7x 節點集區 (np-0 和 np-1)，每個集區包含兩個 tpu7x-standard-4t 節點 (每個節點集區有八個 TPU7x 晶片，採用 2x2x2 拓撲，總共 16 個 TPU7x 晶片)：

    gcloud container node-pools create "${NODE_POOL_NAME_0}" \
    --cluster="${CLUSTER_NAME}" \
    --machine-type="${MACHINE_TYPE}" \
    --placement-policy="${WORKLOAD_POLICY_NAME}" \
    --location="${REGION}" \
    --node-locations="${ZONE}" \
    --project="${PROJECT_ID}" \
    --num-nodes=2 \
    --spot

gcloud container node-pools create "${NODE_POOL_NAME_1}" \
    --cluster="${CLUSTER_NAME}" \
    --machine-type="${MACHINE_TYPE}" \
    --placement-policy="${WORKLOAD_POLICY_NAME}" \
    --location="${REGION}" \
    --node-locations="${ZONE}" \
    --project="${PROJECT_ID}" \
    --num-nodes=2 \
    --spot

    本教學課程使用 Spot 容量 (--spot)。建立節點集區後，使用 Spot 容量的執行個體可能需要一段時間才能建立，且在執行工作負載時可能會遭到先占。您也可以使用預留項目或彈性啟動型 VM 建立節點集區，方法是傳遞下列旗標：

    • 特定預留項目：將 --spot 旗標替換為 --reservation-affinity=specific 和 --reservation=RESERVATION_NAME 旗標，其中 RESERVATION_NAME 是 GKE 建立節點集區時使用的預留項目名稱。如要進一步瞭解 TPU 預留項目，請參閱「關於 Cloud TPU 預留項目」。
    • 彈性啟動型 VM：將 --spot 和 --num-nodes=2 旗標替換為 --flex-start、--reservation-affinity=none、--enable-autoscaling、--num-nodes=0、--min-nodes=0 和 --max-nodes=2 旗標，使用彈性啟動消耗選項佈建節點集區。詳情請參閱「使用 TPU 和彈性啟動型 VM 執行小型批次工作負載」。

    每個節點集區的輸出內容大致如下：

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

  4. 設定 kubectl 憑證，以便與叢集通訊：

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

將模型權重暫存至 Cloud Storage

將模型權重暫存至 Cloud Storage bucket，以便執行工作負載：

  1. 建立 Cloud Storage bucket，並啟用統一 bucket 層級存取權：

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

  2. 將 bucket 的 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"

  3. 建立名為 hf-secret 的 Kubernetes Secret，用於儲存 Hugging Face 存取權杖：

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

  4. 擷取 download_model.yaml 設定：

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

    或者，您也可以選擇建立名為 download_model.yaml 的檔案，並加入下列內容：

    apiVersion: batch/v1
kind: Job
metadata:
  name: download-qwen35
  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: tpu7x
      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.5-397B-A17B-FP8 to Local Staging Volume ==="
          hf download Qwen/Qwen3.5-397B-A17B-FP8 --local-dir /local-ssd/Qwen/Qwen3.5-397B-A17B-FP8
          rm -rf /local-ssd/Qwen/Qwen3.5-397B-A17B-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/Qwen gs://${GCS_BUCKET}/

          echo "=== 4. Cleaning up Local Staging Volume ==="
          rm -rf /local-ssd/Qwen /local-ssd/hf-cache
        resources:
          requests:
            cpu: "64"
            memory: "128Gi"
        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 tpu7x-standard-4t / 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: 450Gi

  5. 將模型權重下載至暫存磁碟區，並轉移至 Cloud Storage bucket：

    envsubst '$GCS_BUCKET' < download_model.yaml | kubectl apply -f -

    這個步驟大約需要 30 到 40 分鐘才能完成。

  6. 工作完成後，請確認所有 163 個 .safetensors shard 都位於 Cloud Storage bucket 中：

    gcloud storage ls "gs://${GCS_BUCKET}/Qwen/Qwen3.5-397B-A17B-FP8/*.safetensors" | wc -l

  7. 刪除已完成的暫存工作：

    kubectl delete job download-qwen35

  8. 等待 Pod 啟動，並監控下載工作的記錄：

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

部署多主機 vLLM TorchTPU 伺服器

工作負載資訊清單 (v7x_qwen3_5_server_multihost_slice.yaml) 定義了 Kubernetes JobSet 資源 (vllm-qwen3-5-397b-pd-disagg-v7x)，其中包含三個複製工作，可協調兩個 2x2x2 TPU7x 節點集區中，分散的預填和解碼服務：

  • 預先填入伺服器 (prefill)：在一個 2x2x2 TPU7x 節點集區中執行兩個副本 (兩個 tpu7x-standard-4t 節點，每部主機有四個雙核心晶片，共 16 個 TPU 核心)。Pod 會形成雙節點 Ray 叢集，並在通訊埠 8000 上啟動 Qwen/Qwen3.5-397B-A17B-FP8 的 vLLM TorchTPU，其中 --data-parallel-size=16、--tensor-parallel-size=1、--enable-expert-parallel、--load-format=runai_streamer 和 TPURaidenConnector 設定為 kv_producer。
  • 解碼伺服器 (decode)：在第二個 2x2x2 TPU7x 節點集區 (共 16 個 TPU 核心) 執行兩個副本，並在 cloud.google.com/gke-nodepool 使用 Pod 親和性和反親和性規則，將預先填入和解碼作業隔離至不同的節點集區。Pod 會形成雙節點 Ray 叢集，並在通訊埠 7200 啟動 vLLM TorchTPU，TPURaidenConnector 則設為 kv_consumer。
  • 分解 Proxy (proxy)：在通訊埠 10000 上執行 toy_proxy_server.py，協調預先填入伺服器 (通訊埠 8000) 和解碼伺服器 (通訊埠 7200) 之間的要求轉送和 KV 快取移轉。

如要部署及監控解構式 vLLM 伺服器，請完成下列步驟：

  1. 在叢集上安裝 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

  2. 擷取伺服器資訊清單：

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

    或者，您也可以建立名為 v7x_qwen3_5_server_multihost_slice.yaml 的檔案，並在其中加入下列內容：

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: vllm-qwen3-5-397b-pd-disagg-v7x
spec:
  failurePolicy:
    restartStrategy: BlockingRecreate
    maxRestarts: 5
  ttlSecondsAfterFinished: 86400
  network:
    enableDNSHostnames: true
  replicatedJobs:
  # ---------------------------------------------------------
  # PREFILL SERVERS (2 Replicas, 1 Node each = 4 dual-core chips / 8 TPU devices per host = 16 TPU devices on 2x2x2 v7x 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: 2x2x2
              cloud.google.com/gke-tpu-accelerator: tpu7x
            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.5-397B-A17B-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": 8}' --block &

                    echo "Waiting for worker node to join Ray cluster (expecting 16.0 TPU total across 2 nodes)..."
                    until ray status 2>/dev/null | grep -E -q "16(\.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.5 397B FP8, DP=16, TP=1, Port 8000) ==="
                    PREFILL_ARGS=(
                      gs://${GCS_BUCKET}/Qwen/Qwen3.5-397B-A17B-FP8
                      --served-model-name=Qwen/Qwen3.5-397B-A17B-FP8
                      --api-server-count=1
                      --load-format=runai_streamer
                      --enable-ep-weight-filter
                      --tensor-parallel-size=1
                      --data-parallel-size=16
                      --data-parallel-size-local=8
                      --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=256
                      --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
                      --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": 8}' --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: ONEHOT_MOE_PERMUTE_THRESHOLD
                value: "2048"
              - 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"
              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 dual-core chips / 8 TPU devices per host = 16 TPU devices on 2x2x2 v7x 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: 2x2x2
              cloud.google.com/gke-tpu-accelerator: tpu7x
            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.5-397B-A17B-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": 8}' --block &

                    echo "Waiting for worker node to join Ray cluster (expecting 16.0 TPU total across 2 nodes)..."
                    until ray status 2>/dev/null | grep -E -q "16(\.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.5 397B FP8, DP=16, TP=1, Port 7200) ==="
                    DECODE_ARGS=(
                      gs://${GCS_BUCKET}/Qwen/Qwen3.5-397B-A17B-FP8
                      --served-model-name=Qwen/Qwen3.5-397B-A17B-FP8
                      --api-server-count=1
                      --load-format=runai_streamer
                      --enable-ep-weight-filter
                      --tensor-parallel-size=1
                      --data-parallel-size=16
                      --data-parallel-size-local=8
                      --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=256
                      --host=0.0.0.0
                      --port=7200
                      --async-scheduling
                      --no-enable-prefix-caching
                      --kv-cache-dtype=fp8
                      --attention-backend=CUSTOM
                      --block-size=8192
                      --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": 8}' --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: ONEHOT_MOE_PERMUTE_THRESHOLD
                value: "2048"
              - 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"
              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" }

  3. 將 Cloud Storage bucket 名稱替換為 v7x_qwen3_5_server_multihost_slice.yaml，然後套用資訊清單：

    envsubst '${GCS_BUCKET} ${VLLM_IMAGE}' < "v7x_qwen3_5_server_multihost_slice.yaml" | kubectl apply -f -

    這項設定會使用一個 2x2x2 TPU7x 切片 (兩個節點) 進行預填，並使用一個 2x2x2 TPU7x 切片 (兩個節點) 進行解碼。如要調整任一集區中的節點數量，請更新 replicas 欄位 (replicas: NUM_NODES)。

  4. 監控預先填入和解碼 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

    首次執行時，初始模型載入和編譯作業約需 30 到 45 分鐘。

    如果您使用 Spot 或彈性啟動型 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.5-397B 模型進行基準測試，請完成下列步驟：

  1. 擷取基準測試資訊清單：

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

    或者，您也可以建立名為 v7x_qwen3_5_benchmark_multihost_slice.yaml 的檔案，並在其中加入下列內容：

    apiVersion: batch/v1
kind: Job
metadata:
  name: vllm-qwen3-5-397b-pd-disagg-v7x-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.5-397B-A17B-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-5-397b-pd-disagg-v7x"

  2. 套用基準測試工作資訊清單，並等待基準測試 Pod 準備就緒：

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

  3. 監控基準測試工作記錄：

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

清除所用資源

如要避免系統向您的 Google Cloud 帳戶收取本教學課程所用資源的費用，請刪除 Kubernetes 工作負載、TPU7x 節點集區、GKE 叢集和 Cloud Storage bucket：

  1. 刪除 Kubernetes 工作負載：

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

  2. 刪除 TPU7x 節點集區、工作負載配置政策和 GKE 叢集：

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

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

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

gcloud compute resource-policies delete "${WORKLOAD_POLICY_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" \
    --quiet || true

  3. (選用) 刪除儲存暫存模型權重的 Cloud Storage bucket：

    gcloud storage rm --recursive "gs://${GCS_BUCKET}" --project="${PROJECT_ID}" --quiet || true

後續步驟