TPU v6e 가속기에서 멀티 호스트 vLLM TorchTPU 및 PD 분리를 사용하여 GKE에서 Qwen3-30B-A3B 서빙

이 튜토리얼에서는 Google Kubernetes Engine(GKE)에서 전문가 혼합 모델인 Qwen3-30B-A3B(Qwen/Qwen3-30B-A3B-FP8)를 배포하고 서빙하는 방법을 보여줍니다. 이 배포에서는 멀티 호스트 TPU v6e(Trillium) 가속기와 TorchTPU(torchtpu) 및 프리필 디코딩(PD) 분리가 포함된 vLLM TPU 서빙 프레임워크를 사용합니다.

이 튜토리얼에 설명된 배포 구성에서는 각각 2x4 토폴로지로 정렬된 멀티 호스트 ct6e-standard-4t 노드 풀 2개를 프로비저닝합니다(노드 풀당 노드 2개에 TPU v6e 칩 8개, 총 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. 영구 디스크 SSD(GB)로 필터링하고 대상 리전에 사용 가능한 SSD 할당량이 충분한지 확인합니다.

    3. 사용 중인 리전 외부 IPv4 주소로 필터링하고 대상 리전에 IP 주소 할당량이 충분한지 확인합니다.

Hugging Face 액세스 토큰 가져오기

이 튜토리얼을 사용하려면 Hugging Face 액세스 토큰이 필요합니다. Hugging Face에서 무료 계정에 가입할 수 있습니다. 계정이 있으면 액세스 토큰을 생성합니다.

  1. Hugging Face에 오신 것을 환영합니다 페이지에서 계정 아바타를 클릭하고 액세스 토큰을 선택합니다.
  2. 액세스 토큰 페이지에서 새 토큰 만들기를 클릭합니다.
  3. 읽기 토큰 유형을 선택하고 토큰 이름을 입력합니다.
  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="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 클러스터 만들기

일반 모드에서 GKE 클러스터를 만들고 프리필 및 디코딩 분리를 위해 멀티 호스트 TPU v6e 노드 풀 2개를 연결하려면 다음 단계를 완료하세요.

  1. 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. 각각 ct6e-standard-4t 노드 2개를 포함하는 멀티 호스트 TPU v6e 노드 풀 2개를 만듭니다(노드 풀당 2x4 토폴로지에 TPU v6e 칩 8개, 총 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)을 사용합니다. 스팟 용량을 사용하는 인스턴스는 노드 풀이 생성된 후 생성되는 데 시간이 걸릴 수 있으며 워크로드 실행 중에 선점될 수 있습니다. 다음 플래그를 전달하여 예약 또는 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

  3. 클러스터와 통신하도록 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. Hugging Face 액세스 토큰을 저장할 hf-secret이라는 Kubernetes 보안 비밀을 만듭니다.

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

  4. 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

  5. Qwen/Qwen3-30B-A3B-FP8 모델 가중치를 스테이징 볼륨에 다운로드하고 Cloud Storage 버킷으로 전송합니다.

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

    이 단계를 완료하는 데 약 5~10분이 걸립니다.

  6. 포드가 시작될 때까지 기다린 후 다운로드 작업 로그를 모니터링합니다.

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

  7. 작업이 완료되면 16개의 .safetensors 샤드가 모두 Cloud Storage 버킷에 있는지 확인합니다.

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

  8. 완료된 스테이징 작업을 삭제합니다.

    kubectl delete job download-qwen3

멀티 호스트 vLLM TorchTPU 서버 배포

워크로드 매니페스트(v6e_qwen3_server_multihost_slice.yaml)는 두 개의 2x4 TPU v6e 노드 풀에서 분리된 프리필 및 디코딩 서빙을 조정하는 세 개의 복제된 작업이 있는 Kubernetes JobSet 리소스(vllm-qwen3-30b-pd-disagg-v6e)를 정의합니다.

  • 프리필 서버(prefill): 하나의 2x4 TPU v6e 노드 풀에서 두 개의 복제본을 실행합니다(호스트당 4개의 칩이 있는 ct6e-standard-4t 노드 2개, 총 8개의 TPU 코어). 포드는 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): 두 번째 2x4 TPU v6e 노드 풀에서 복제본 2개를 실행합니다(총 TPU 코어 8개). JobSet는 cloud.google.com/gke-nodepool 라벨의 포드 어피니티 및 안티어피니티 규칙을 사용하여 프리필 및 디코딩을 별도의 노드 풀로 격리합니다. 포드는 2노드 Ray 클러스터를 형성하고 TPURaidenConnector가 kv_consumer로 구성된 포트 7200에서 vLLM TorchTPU를 실행합니다.

  • 분리 프록시(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/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" }

  3. Cloud Storage 버킷 이름과 컨테이너 이미지를 v6e_qwen3_server_multihost_slice.yaml 파일에 대체하고 매니페스트를 적용합니다.

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

    이 구성은 프리필에 2x4 TPU v6e 슬라이스(노드 2개) 1개와 디코딩에 2x4 TPU v6e 슬라이스(노드 2개) 1개를 사용합니다. 풀의 노드 수를 확장하려면 replicas 필드(replicas: NUM_NODES)를 업데이트합니다. NUM_NODES를 사용하려는 노드 수로 바꿉니다.

  4. 애플리케이션 시작이 완료될 때까지 프리필 및 디코딩 포드 로그를 모니터링합니다.

    # 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분이 걸립니다.

    스팟 용량 또는 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 모델을 분리 프록시에 대해 벤치마킹하려면 다음 단계를 완료합니다.

  1. 벤치마크 매니페스트를 가져옵니다.

    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"

  2. 벤치마크 작업 매니페스트를 적용합니다.

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

  3. 벤치마크 포드가 준비될 때까지 기다린 후 벤치마크 작업 로그를 모니터링합니다.

    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

  4. 벤치마크가 완료되면 다음과 같은 측정항목이 포함된 출력을 찾습니다.

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

삭제

이 튜토리얼에서 사용한 리소스 비용이 Google Cloud 계정에 청구되지 않도록 하려면 Kubernetes 워크로드, TPU v6e 노드 풀, GKE 클러스터, Cloud Storage 버킷을 삭제합니다.

  1. Kubernetes 워크로드를 삭제합니다.

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

  2. 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

  3. 스테이징된 모델 가중치를 저장하는 Cloud Storage 버킷을 삭제합니다. 모델 가중치를 유지하려면 이 단계를 건너뛰거나 원하는 위치에 복사합니다.

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

다음 단계