O Google usa tecnologia de IA na tradução de conteúdos para seu idioma de preferência. As traduções com IA podem ter erros.

Disponibilizar o Qwen3-30B-A3B no GKE com vLLM TorchTPU de vários hosts e desagregação DP em aceleradores de TPU v6e

Neste tutorial, mostramos como implantar e disponibilizar o Qwen3-30B-A3B (Qwen/Qwen3-30B-A3B-FP8), um modelo de mix de especialistas, no Google Kubernetes Engine (GKE). A implantação usa aceleradores TPU v6e (Trillium) de vários hosts e o framework de veiculação de TPU vLLM com TorchTPU (torchtpu) e desagregação de pré-preenchimento e decodificação (DP).

Na configuração de implantação descrita neste tutorial, você provisiona dois pools de nós ct6e-standard-4t de vários hosts, cada um organizado em uma topologia 2x4 (oito chips de TPU v6e em dois nós por pool de nós, totalizando 16 chips de TPU v6e). Use a API JobSet do Kubernetes e o Ray para fazer o seguinte:

Execute servidores de pré-preenchimento e decodificação desagregados de vários hosts em frações separadas da TPU v6e com o paralelismo de dados e o paralelismo de especialistas ativados.
Encaminhar solicitações por um proxy de desagregação.

Para executar um fluxo de trabalho de disponibilização de TorchTPU vLLM de vários hosts semelhante em aceleradores TPU7x (Ironwood) com o modelo Qwen3.5-397B maior, consulte Disponibilizar o Qwen3.5-397B no GKE com TorchTPU vLLM de vários hosts em aceleradores TPU7x.

Objetivos

Configure o cluster do GKE e os pools de nós de TPU v6e de vários hosts.
Armazene os pesos do modelo em um bucket do Cloud Storage.
Implante os servidores de preenchimento, decodificação e proxy do vLLM TorchTPU de vários hosts com o JobSet.
Executar uma carga de trabalho de comparativo de mercado de exibição e analisar os resultados.
Fazer a limpeza.

Custos

Neste documento, você usará os seguintes componentes faturáveis do Google Cloud:

Para gerar uma estimativa de custo baseada na projeção de uso deste tutorial, use a calculadora de preços.

Ao concluir as tarefas descritas neste documento, é possível evitar o faturamento contínuo excluindo os recursos criados. Para mais informações, consulte Limpeza.

Antes de começar

Configure o ambiente para o tutorial:

Ative a API Kubernetes Engine no projeto do Google Cloud :

Ativar a API Kubernetes Engine
Instale e inicialize a Google Cloud CLI e instale kubectl. Se você usa o Cloud Shell, as duas ferramentas já estão pré-instaladas.
Verifique se o projeto tem cota de SSD e cota de endereço IP externo suficientes na região em que a capacidade da TPU está localizada:
1. No console do Google Cloud , acesse a página Cotas e limites do sistema.
  
  Acesse Cotas e limites do sistema
2. Filtre por SSD de disco permanente (GB) e verifique se você tem cota de SSD suficiente disponível na região de destino.
3. Filtre por Endereços IPv4 externos regionais em uso e verifique se você tem cota de endereço IP suficiente na região de destino.

Receber um token de acesso do Hugging Face

Você precisa de um token de acesso do Hugging Face para usar este tutorial. Você pode se inscrever para uma conta sem custo financeiro no Hugging Face. Depois de criar uma conta, gere um token de acesso:

Na página Welcome to Hugging Face, clique no avatar da sua conta e selecione Tokens de acesso.
Na página Tokens de acesso, clique em Criar novo token.
Selecione o tipo de token Read e insira um nome para ele.
Seu token de acesso será exibido. Salve o token em um local seguro.

Funções exigidas

Para conseguir as permissões necessárias para concluir o tutorial, peça ao administrador para conceder a você os seguintes papéis do IAM no seu projeto:

Administrador do Kubernetes Engine (roles/container.admin)
Administrador do Compute (roles/compute.admin)
Administrador de armazenamento (roles/storage.admin)
Usuário da conta de serviço (roles/iam.serviceAccountUser)

Para mais informações sobre a concessão de papéis, consulte Gerenciar o acesso a projetos, pastas e organizações.

Também é possível conseguir as permissões necessárias usando papéis personalizados ou outros papéis predefinidos.

Configurar as variáveis de ambiente

Para configurar as variáveis de ambiente, siga estas etapas:

Configure sua autenticação:

gcloud auth login
gcloud auth application-default login

Exporte as variáveis de ambiente que você vai usar neste tutorial:
```
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}"
```
Substitua:
- YOUR_PROJECT_ID: o ID do seu projeto do Google Cloud .
- YOUR_REGION: a região em que você quer implantar o cluster, como us-central1.
- YOUR_ZONE: a zona em que você quer implantar os pools de nós da TPU v6e, como us-central1-c.
- YOUR_CLUSTER_NAME: o nome do cluster do GKE.
- YOUR_RESERVATION_NAME: se você estiver usando uma reserva, o nome da sua reserva de TPU v6e.
- YOUR_GCS_BUCKET: o nome do bucket do Cloud Storage que você usa para preparar os pesos do modelo.
- YOUR_HF_TOKEN: seu token de acesso do Hugging Face.
Se você estiver usando uma imagem de contêiner personalizada do vLLM TorchTPU, substitua o URI de imagem padrão vllm/vllm-torchtpu na variável VLLM_IMAGE pelo URI da sua imagem. Para imagens públicas, consulte o repositório vLLM TorchTPU.

Criar o cluster do GKE

Para criar um cluster do GKE no modo Standard e anexar dois pools de nós de TPU v6e de vários hosts para desagregação de pré-preenchimento e decodificação, siga estas etapas:

Crie um cluster regional do GKE com a federação de identidade da carga de trabalho para GKE ativada:

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"

O resultado será o seguinte:

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

Crie dois pools de nós de TPU v6e de vários hosts, cada um contendo dois nós ct6e-standard-4t (oito chips de TPU v6e em uma topologia 2x4 por pool de nós, totalizando 16 chips de 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

Neste tutorial, usamos a capacidade spot (--spot). As instâncias que usam essa capacidade podem levar algum tempo para serem criadas após a criação do pool de nós e podem ser interrompidas durante as execuções da carga de trabalho. Também é possível criar os pools de nós usando uma reserva ou VMs flex-start transmitindo as seguintes flags:

Reserva específica: crie o pool de nós usando uma reserva específica:

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

Para saber mais sobre reservas de TPU, consulte Sobre as reservas do Cloud TPU.

VMs flex-start: provisione o pool de nós usando a opção de consumo de início flexível:

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

Para mais informações, consulte Executar uma carga de trabalho em lote pequena com TPUs e VMs de início flexível.

A saída de cada pool de nós é semelhante a esta:

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

Configure as credenciais de kubectl para se comunicar com o cluster:

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

Organizar os pesos do modelo no Cloud Storage

Prepare os pesos do modelo em um bucket do Cloud Storage para as execuções da carga de trabalho:

Crie um bucket do Cloud Storage com acesso uniforme no nível do bucket:

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

Conceda permissões do Identity and Access Management (IAM) no bucket à conta de serviço padrão do Compute Engine e à conta de serviço padrão do 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"

Crie um secret do Kubernetes chamado hf-secret para armazenar o token de acesso do Hugging Face:

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

Busque o manifesto download_qwen3_30b.yaml:

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

Como alternativa, crie um arquivo chamado download_qwen3_30b.yaml com o seguinte conteúdo:

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

Faça o download dos pesos do modelo Qwen/Qwen3-30B-A3B-FP8 para o volume de preparo e transfira-os para o bucket do Cloud Storage:
```
envsubst '$GCS_BUCKET' < download_qwen3_30b.yaml | kubectl apply -f -
```
Essa etapa leva de 5 a 10 minutos para ser concluída.

Aguarde o pod iniciar e monitore os registros do job de download:

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

Depois que o job for concluído, verifique se todos os 16 fragmentos .safetensors estão no bucket do Cloud Storage:
```
gcloud storage ls "gs://${GCS_BUCKET}/Qwen/Qwen3-30B-A3B-FP8/*.safetensors" | wc -l
```
Exclua o job de preparo concluído:
```
kubectl delete job download-qwen3
```

Implantar o servidor TorchTPU do vLLM de vários hosts

O manifesto da carga de trabalho (v6e_qwen3_server_multihost_slice.yaml) define um recurso JobSet do Kubernetes (vllm-qwen3-30b-pd-disagg-v6e) com três jobs replicados que coordenam a disponibilização de pré-preenchimento e decodificação desagregados nos dois pools de nós de TPU v6e 2x4:

Servidor de preenchimento prévio (prefill): executa duas réplicas em um pool de nós de TPU v6e 2x4 (dois nós ct6e-standard-4t com quatro chips por host, totalizando oito núcleos de TPU). Os pods formam um cluster do Ray de dois nós e iniciam o vLLM TorchTPU na porta 8000 para Qwen/Qwen3-30B-A3B-FP8. O servidor usa TPURaidenConnector configurado como kv_producer com as seguintes flags:
- --data-parallel-size=8
- --tensor-parallel-size=1
- --enable-expert-parallel
- --load-format=runai_streamer
Servidor de decodificação (decode): executa duas réplicas no segundo pool de nós de TPU v6e 2x4 (oito núcleos de TPU no total). O JobSet usa regras de afinidade e antiafinidade de pods no rótulo cloud.google.com/gke-nodepool para isolar o preenchimento prévio e a decodificação em pools de nós separados. Os pods formam um cluster do Ray de dois nós e iniciam o vLLM TorchTPU na porta 7200 com TPURaidenConnector configurado como kv_consumer.
Proxy de desagregação (proxy): executa o script toy_proxy_server.py na porta 10000 para coordenar o roteamento de solicitações e a transferência de cache de KV entre o servidor de preenchimento prévio (porta 8000) e o servidor de decodificação (porta 7200).

Para implantar e monitorar os servidores vLLM desagregados, siga estas etapas:

Instale o CustomResourceDefinition e o controlador do JobSet do Kubernetes no cluster:

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

Busque o manifesto do servidor:

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

Como alternativa, crie um arquivo chamado v6e_qwen3_server_multihost_slice.yaml com o seguinte conteúdo:

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" }

Substitua o nome do bucket do Cloud Storage e a imagem do contêiner no arquivo v6e_qwen3_server_multihost_slice.yaml e aplique o manifesto:
```
envsubst '${GCS_BUCKET} ${VLLM_IMAGE}' < "v6e_qwen3_server_multihost_slice.yaml" | kubectl apply -f -
```
Essa configuração usa uma fração de TPU v6e 2x4 (dois nós) para pré-preenchimento e uma fração de TPU v6e 2x4 (dois nós) para decodificação. Para escalonar o número de nós em qualquer pool, atualize o campo replicas (replicas: NUM_NODES). Substitua NUM_NODES pelo número de nós que você quer usar.
Monitore os registros de pod de preenchimento prévio e decodificação até que a inicialização do aplicativo seja concluída:
```
# 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
```
O carregamento e a compilação iniciais do modelo levam de 15 a 25 minutos na primeira execução.

Se você estiver usando capacidade preemptiva, como capacidade spot ou VMs flex-start, pode levar mais tempo após a criação do pool de nós para que os nós de fração de TPU sejam criados, e eles podem ser interrompidos durante as execuções de carga de trabalho. Para mais detalhes, consulte Instâncias de VM preemptivas. Para verificar se os nós de fração de TPU foram criados, consulte Verificar o status e as condições do nó.

Aguarde até que os registros do servidor de preenchimento prévio e decodificação informem Application startup complete antes de executar a carga de trabalho de comparativo de mercado:
```
(APIServer pid=1) INFO:     Started server process [1]
(APIServer pid=1) INFO:     Waiting for application startup.
(APIServer pid=1) INFO:     Application startup complete.
```

Executar uma carga de trabalho de comparativo de mercado

Para comparar o modelo Qwen3-30B-A3B implantado usando uma carga de trabalho de entrada de 8.192 tokens e saída de 1.024 tokens (8k/1k) com o proxy de desagregação, siga estas etapas:

Busque o manifesto de comparativo de mercado:

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

Como alternativa, crie um arquivo chamado v6e_qwen3_benchmark_multihost_slice.yaml com o seguinte conteúdo:

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"

Aplique o manifesto do job de comparativo de mercado:

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

Aguarde até que o pod de comparativo de mercado esteja pronto e monitore os registros do job de comparativo de mercado:

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

Depois que o comparativo de mercado for concluído, procure uma saída que contenha métricas como as seguintes:
- Output token throughput
- Total token throughput
- TTFT
- TPOT

Limpar

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados neste tutorial, exclua as cargas de trabalho do Kubernetes, os pools de nós da TPU v6e, o cluster do GKE e o bucket do Cloud Storage:

Exclua as cargas de trabalho do Kubernetes:

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

Exclua os pools de nós da TPU v6e e o cluster do 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

Exclua o bucket do Cloud Storage que armazena os pesos do modelo em etapas. Se você quiser manter os pesos do modelo, pule esta etapa ou copie-os para o local de sua preferência:
```
gcloud storage rm --recursive "gs://${GCS_BUCKET}" --project="${PROJECT_ID}" --quiet || true
```

A seguir

Saiba mais sobre o vLLM no Cloud TPU.
Saiba mais sobre a TPU v6e (Trillium).
Saiba mais sobre o Cloud TPU.