在 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 拓扑（每个节点池包含两个节点，每个节点有 8 个 TPU7x 芯片，总共 16 个 TPU7x 芯片）。您可以使用 Kubernetes JobSet API 和 Ray 在单独的 TPU7x 切片上运行解耦的多主机预填充和解码服务器，并启用数据并行和专家并行，然后通过解耦代理路由请求。

目标

  • 设置 GKE 集群和多主机 TPU7x 节点池。
  • 将模型权重暂存在 Cloud Storage 存储桶中。
  • 使用 JobSet 部署多主机 vLLM TorchTPU 预填充、解码和代理服务器。
  • 运行服务基准工作负载并分析结果。
  • 清理。

费用

在本文档中，您将使用Google Cloud的以下收费组件：

如需根据您的预计使用量来估算费用，请使用价格计算器。

完成本文档中描述的任务后，您可以通过删除所创建的资源来避免继续计费。如需了解详情，请参阅清理。

准备工作

  • 您需要拥有 Hugging Face 访问令牌才能使用本教程。您可以在 Hugging Face 注册免费账号。拥有账号后，生成访问令牌：

    1. 在欢迎使用 Hugging Face 页面上，点击您的账号头像，然后选择访问令牌。
    2. 在访问令牌页面上，点击创建新令牌。
    3. 选择读取令牌类型，然后输入令牌名称。
    4. 系统会显示您的访问令牌。将令牌保存在安全的位置。

  • 在 Google Cloud 项目中启用 Kubernetes Engine API：

    启用 Kubernetes Engine API

  • 安装并初始化 Google Cloud CLI，然后安装 kubectl。如果您使用 Cloud Shell，则这两个工具均已预安装。

  • 确保您的项目在 TPU 容量所在的区域中有足够的 SSD 配额和外部 IP 地址配额：

    1. 在 Google Cloud 控制台中，前往配额和系统限制页面。

      进入“配额和系统限制”

    2. 按永久性磁盘 SSD (GB) 过滤，并验证您在目标区域中是否有足够的可用 SSD 配额。

    3. 按正在使用的区域级外部 IPv4 地址过滤，并验证目标区域中是否有足够的 IP 地址配额。

如需获得完成本教程所需的权限，请让您的管理员为您授予项目的以下 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 存储桶的名称。
    • YOUR_HF_TOKEN：您的 Hugging Face 访问令牌。
    • YOUR_VLLM_IMAGE：要使用的 vLLM TorchTPU 映像。如需了解公开映像，请参阅 vLLM TorchTPU 代码库。

创建 GKE 集群

如需在 Standard 模式下创建 GKE 集群，并挂接两个多主机 TPU7x 节点池以实现预填充和解码解耦，请完成以下步骤：

  1. 创建启用了 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-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 节点（每个节点池在 2x2x2 拓扑中包含 8 个 TPU7x 芯片，总共 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 容量的实例在节点池创建后可能需要一些时间才能创建，并且可能会在工作负载运行期间被抢占。您还可以通过传递以下标志，使用预留或灵活启动型虚拟机创建节点池：

    • 特定预留：将 --spot 标志替换为 --reservation-affinity=specific 和 --reservation=RESERVATION_NAME 标志，其中 RESERVATION_NAME 是 GKE 在创建节点池时使用的预留的名称。如需详细了解 TPU 预留，请参阅 Cloud TPU 预留简介。
    • 灵活启动型虚拟机：将 --spot 和 --num-nodes=2 标志替换为 --flex-start、--reservation-affinity=none、--enable-autoscaling、--num-nodes=0、--min-nodes=0 和 --max-nodes=2 标志，以使用灵活启动使用选项预配节点池。如需了解详情，请参阅使用 TPU 和灵活启动虚拟机运行小型批量工作负载。

    每个节点池的输出类似于以下内容：

    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 存储桶中暂存模型权重，以便工作负载运行：

  1. 创建具有统一存储桶级访问权限的 Cloud Storage 存储桶：

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

  2. 向默认 Compute Engine 服务账号和默认 Kubernetes 服务账号授予存储桶的 Identity and Access Management (IAM) 权限：

    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 存储桶：

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

    此步骤大约需要 30 到 40 分钟才能完成。

  6. 作业完成后，验证所有 163 个 .safetensors 分片是否都在 Cloud Storage 存储桶中：

    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)：在端口 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 存储桶名称替换为 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 虚拟机或灵活启动型虚拟机），则在创建节点池后，TPU 切片节点可能需要更长时间才能创建，并且在工作负载运行期间可能会被抢占。如需了解详情，请参阅抢占式虚拟机实例。如需检查 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 存储桶：

  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 存储桶：

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

后续步骤