Google utilise la technologie IA pour traduire le contenu dans votre langue préférée. Les traductions générées par IA peuvent contenir des erreurs.

Mettre en service Qwen3-30B-A3B sur GKE avec vLLM TorchTPU multi-hôte et désagrégation de disque persistant sur des accélérateurs TPU v6e

Ce tutoriel vous explique comment déployer et diffuser Qwen3-30B-A3B (Qwen/Qwen3-30B-A3B-FP8), un modèle Mixture of Experts, sur Google Kubernetes Engine (GKE). Le déploiement utilise des accélérateurs TPU v6e (Trillium) multi-hôtes et le framework de diffusion TPU vLLM avec TorchTPU (torchtpu) et la désagrégation de préremplissage-décodage (PD).

Dans la configuration de déploiement décrite dans ce tutoriel, vous provisionnez deux pools de nœuds ct6e-standard-4t multi-hôtes, chacun organisé selon une topologie 2x4 (huit puces TPU v6e sur deux nœuds par pool de nœuds, pour un total de 16 puces TPU v6e). Vous utilisez l'API Kubernetes JobSet et Ray pour effectuer les opérations suivantes :

Exécutez des serveurs de préremplissage et de décodage multi-hôtes désagrégés sur des tranches TPU v6e distinctes avec le parallélisme des données et le parallélisme des experts activés.
Acheminez les requêtes via un proxy de désagrégation.

Pour exécuter un workflow de diffusion TorchTPU vLLM multi-hôte similaire sur des accélérateurs TPU7x (Ironwood) avec le modèle Qwen3.5-397B plus volumineux, consultez Diffuser Qwen3.5-397B sur GKE avec TorchTPU vLLM multi-hôte sur des accélérateurs TPU7x.

Objectifs

Configurez votre cluster GKE et vos pools de nœuds TPU v6e multi-hôtes.
Préparer les pondérations du modèle dans un bucket Cloud Storage.
Déployez les serveurs de préremplissage, de décodage et de serveur proxy vLLM TorchTPU multi-hôtes avec JobSet.
Exécuter une charge de travail de benchmark de diffusion et analyser les résultats
Effectuer un nettoyage.

Coûts

Dans ce document, vous utilisez les composants facturables suivants deGoogle Cloud :

Obtenez une estimation des coûts en fonction de votre utilisation prévue à l'aide du simulateur de coût.

Une fois que vous avez terminé les tâches décrites dans ce document, vous pouvez éviter de continuer à payer des frais en supprimant les ressources que vous avez créées. Pour en savoir plus, consultez la section Effectuer un nettoyage.

Avant de commencer

Configurez votre environnement pour le tutoriel :

Activez l'API Kubernetes Engine dans votre projet Google Cloud :

Activer l'API Kubernetes Engine
Installez et initialisez la Google Cloud CLI, puis installez kubectl. Si vous utilisez Cloud Shell, les deux outils sont préinstallés.
Assurez-vous que votre projet dispose d'un quota de SSD et d'un quota d'adresses IP externes suffisants dans la région où se trouve votre capacité TPU :
1. Dans la console Google Cloud , accédez à la page Quotas et limites du système.
  
  Accéder à la page Quotas et limites du système
2. Filtrez par Disque persistant SSD (Go) et vérifiez que vous disposez d'un quota de SSD suffisant dans votre région cible.
3. Filtrez par Adresses IPv4 externes régionales en cours d'utilisation et vérifiez que vous disposez d'un quota d'adresses IP suffisant dans votre région cible.

Obtenir un jeton d'accès Hugging Face

Vous avez besoin d'un jeton d'accès Hugging Face pour suivre ce tutoriel. Vous pouvez créer un compte sans frais sur Hugging Face. Une fois que vous disposez d'un compte, générez un jeton d'accès :

Sur la page Welcome to Hugging Face (Bienvenue sur Hugging Face), cliquez sur l'avatar de votre compte, puis sélectionnez Access tokens (Jetons d'accès).
Sur la page Access tokens (Jetons d'accès), cliquez sur Create new token (Créer un jeton).
Sélectionnez le type de jeton Read (Lecture) et saisissez un nom pour votre jeton.
Votre jeton d'accès s'affiche. Enregistrez le jeton dans un endroit sûr.

Rôles requis

Pour obtenir les autorisations nécessaires pour suivre ce tutoriel, demandez à votre administrateur de vous accorder les rôles IAM suivants sur votre projet :

Administrateur Kubernetes Engine (roles/container.admin)
Administrateur de Compute (roles/compute.admin)
Administrateur de l'espace de stockage (roles/storage.admin)
Utilisateur du compte de service (roles/iam.serviceAccountUser)

Pour en savoir plus sur l'attribution de rôles, consultez la page Gérer l'accès aux projets, aux dossiers et aux organisations.

Vous pouvez également obtenir les autorisations requises avec des rôles personnalisés ou d'autres rôles prédéfinis.

Configurer vos variables d'environnement

Pour configurer vos variables d'environnement, procédez comme suit :

Configurez votre authentification :

gcloud auth login
gcloud auth application-default login

Exportez les variables d'environnement que vous utiliserez tout au long de ce tutoriel :
```
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}"
```
Remplacez les éléments suivants :
- YOUR_PROJECT_ID : ID de votre projetGoogle Cloud .
- YOUR_REGION : région dans laquelle vous souhaitez déployer votre cluster, par exemple us-central1.
- YOUR_ZONE : zone dans laquelle vous souhaitez déployer vos pools de nœuds TPU v6e, par exemple us-central1-c.
- YOUR_CLUSTER_NAME : nom de votre cluster GKE.
- YOUR_RESERVATION_NAME : si vous utilisez une réservation, nom de votre réservation de TPU v6e.
- YOUR_GCS_BUCKET : nom du bucket Cloud Storage que vous utilisez pour stocker les pondérations du modèle.
- YOUR_HF_TOKEN : votre jeton d'accès Hugging Face.
Si vous utilisez une image de conteneur vLLM TorchTPU personnalisée, remplacez l'URI d'image vllm/vllm-torchtpu par défaut dans la variable VLLM_IMAGE par l'URI de votre image. Pour les images publiques, consultez le dépôt vLLM TorchTPU.

Créer le cluster GKE

Pour créer un cluster GKE en mode Standard et associer deux pools de nœuds TPU v6e multi-hôtes pour la désagrégation du préremplissage et du décodage, procédez comme suit :

Créez un cluster GKE régional avec la fédération d'Workload Identity pour GKE activée :

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"

Le résultat ressemble à ce qui suit :

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

Créez deux pools de nœuds TPU v6e multi-hôtes, chacun contenant deux nœuds ct6e-standard-4t (huit puces TPU v6e dans une topologie 2x4 par pool de nœuds, pour un total de 16 puces 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

Ce tutoriel utilise la capacité Spot (--spot). La création d'instances utilisant la capacité Spot peut prendre du temps après la création du pool de nœuds, et elles peuvent être préemptées pendant l'exécution de la charge de travail. Vous pouvez également créer les pools de nœuds à l'aide d'une réservation ou de VM à démarrage flexible en transmettant les options suivantes :

Réservation spécifique : créez le pool de nœuds à l'aide d'une réservation spécifique :

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

Pour en savoir plus sur les réservations TPU, consultez À propos des réservations Cloud TPU.

VM à démarrage flexible : provisionnez le pool de nœuds à l'aide de l'option de consommation Démarrage flexible :

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

Pour en savoir plus, consultez Exécuter une petite charge de travail par lot avec des TPU et des VM à démarrage flexible.

Le résultat de chaque pool de nœuds ressemble à ce qui suit :

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

Configurez les identifiants kubectl pour communiquer avec votre cluster :

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

Stocker les pondérations du modèle dans Cloud Storage

Préparez les pondérations du modèle dans un bucket Cloud Storage pour les exécutions de votre charge de travail :

Créez un bucket Cloud Storage avec un accès uniforme au niveau du bucket :

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

Accordez des autorisations IAM (Identity and Access Management) sur le bucket à votre compte de service Compute Engine par défaut et au compte de service Kubernetes par défaut :

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"

Créez un secret Kubernetes nommé hf-secret pour stocker votre jeton d'accès Hugging Face :

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

Récupérez le fichier manifeste download_qwen3_30b.yaml :

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

Vous pouvez également créer un fichier nommé download_qwen3_30b.yaml avec le contenu suivant :

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

Téléchargez les pondérations du modèle Qwen/Qwen3-30B-A3B-FP8 sur le volume de préproduction et transférez-les vers votre bucket Cloud Storage :
```
envsubst '$GCS_BUCKET' < download_qwen3_30b.yaml | kubectl apply -f -
```
Cette étape prend environ 5 à 10 minutes.

Attendez que le pod démarre et surveillez les journaux du job de téléchargement :

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

Une fois le job terminé, vérifiez que les 16 segments .safetensors se trouvent dans votre bucket Cloud Storage :
```
gcloud storage ls "gs://${GCS_BUCKET}/Qwen/Qwen3-30B-A3B-FP8/*.safetensors" | wc -l
```
Supprimez le job de préproduction terminé :
```
kubectl delete job download-qwen3
```

Déployer le serveur vLLM TorchTPU multi-hôte

Le fichier manifeste de la charge de travail (v6e_qwen3_server_multihost_slice.yaml) définit une ressource JobSet Kubernetes (vllm-qwen3-30b-pd-disagg-v6e) avec trois jobs répliqués qui coordonnent la mise en service désagrégée du préremplissage et du décodage sur les deux pools de nœuds TPU v6e 2x4 :

Serveur de préremplissage (prefill) : exécute deux instances répliquées sur un pool de nœuds TPU v6e 2x4 (deux nœuds ct6e-standard-4t avec quatre puces par hôte, pour un total de huit cœurs TPU). Les pods forment un cluster Ray à deux nœuds et lancent vLLM TorchTPU sur le port 8000 pour Qwen/Qwen3-30B-A3B-FP8. Le serveur utilise TPURaidenConnector configuré en tant que kv_producer avec les options suivantes :
- --data-parallel-size=8
- --tensor-parallel-size=1
- --enable-expert-parallel
- --load-format=runai_streamer
Serveur de décodage (decode) : exécute deux instances répliquées sur le deuxième pool de nœuds TPU v6e 2x4 (huit cœurs TPU au total). Le JobSet utilise des règles d'affinité et d'anti-affinité de pod sur l'étiquette cloud.google.com/gke-nodepool pour isoler le préremplissage et le décodage sur des pools de nœuds distincts. Les pods forment un cluster Ray à deux nœuds et lancent vLLM TorchTPU sur le port 7200 avec TPURaidenConnector configuré sur kv_consumer.
Proxy de désagrégation (proxy) : exécute le script toy_proxy_server.py sur le port 10000 pour coordonner le routage des requêtes et le transfert du cache KV entre le serveur de préremplissage (port 8000) et le serveur de décodage (port 7200).

Pour déployer et surveiller les serveurs vLLM désagrégés, procédez comme suit :

Installez la ressource CustomResourceDefinition et le contrôleur JobSet Kubernetes sur votre 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

Récupérez le fichier manifeste du serveur :

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

Vous pouvez également créer un fichier nommé v6e_qwen3_server_multihost_slice.yaml avec le contenu suivant :

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

Remplacez le nom de votre bucket Cloud Storage et l'image de conteneur dans le fichier v6e_qwen3_server_multihost_slice.yaml, puis appliquez le fichier manifeste :
```
envsubst '${GCS_BUCKET} ${VLLM_IMAGE}' < "v6e_qwen3_server_multihost_slice.yaml" | kubectl apply -f -
```
Cette configuration utilise une tranche TPU v6e 2x4 (deux nœuds) pour le préremplissage et une tranche TPU v6e 2x4 (deux nœuds) pour le décodage. Pour faire évoluer le nombre de nœuds dans l'un ou l'autre des pools, mettez à jour le champ replicas (replicas: NUM_NODES). Remplacez NUM_NODES par le nombre de nœuds que vous souhaitez utiliser.
Surveillez les journaux des pods de préremplissage et de décodage jusqu'à ce que l'application ait démarré :
```
# 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
```
Le chargement et la compilation initiaux du modèle prennent environ 15 à 25 minutes lors de la première exécution.

Si vous utilisez une capacité préemptive telle que la capacité Spot ou des VM Flex-start, la création de vos nœuds de tranche TPU peut prendre plus de temps après la création du pool de nœuds, et ils peuvent être préemptés pendant l'exécution des charges de travail. Pour en savoir plus, consultez Instances de VM préemptives. Pour vérifier si vos nœuds de tranche TPU ont été créés, consultez Vérifier l'état et les conditions du nœud.

Attendez que les journaux des serveurs de préremplissage et de décodage indiquent Application startup complete avant d'exécuter la charge de travail de benchmark :
```
(APIServer pid=1) INFO:     Started server process [1]
(APIServer pid=1) INFO:     Waiting for application startup.
(APIServer pid=1) INFO:     Application startup complete.
```

Exécuter une charge de travail de benchmark

Pour comparer le modèle Qwen3-30B-A3B déployé à l'aide d'une charge de travail d'entrée de 8 192 jetons et de sortie de 1 024 jetons (8k/1k) par rapport au proxy de désagrégation, procédez comme suit :

Récupérez le fichier manifeste du benchmark :

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

Vous pouvez également créer un fichier nommé v6e_qwen3_benchmark_multihost_slice.yaml avec le contenu suivant :

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"

Appliquez le fichier manifeste du job de benchmark :

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

Attendez que le pod de benchmark soit prêt, puis surveillez les journaux du job de benchmark :

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

Une fois le benchmark terminé, recherchez un résultat contenant des métriques telles que les suivantes :
- Output token throughput
- Total token throughput
- TTFT
- TPOT

Effectuer un nettoyage

Pour éviter que les ressources utilisées dans ce tutoriel ne soient facturées sur votre compte Google Cloud , supprimez les charges de travail Kubernetes, les pools de nœuds TPU v6e, le cluster GKE et votre bucket Cloud Storage :

Supprimez les charges de travail Kubernetes :

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

Supprimez les pools de nœuds TPU v6e et le cluster 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

Supprimez le bucket Cloud Storage qui stocke les pondérations du modèle préproduit. Si vous souhaitez conserver les pondérations du modèle, ignorez cette étape ou copiez-les à l'emplacement de votre choix :
```
gcloud storage rm --recursive "gs://${GCS_BUCKET}" --project="${PROJECT_ID}" --quiet || true
```

Étapes suivantes

En savoir plus sur vLLM sur Cloud TPU.
En savoir plus sur les TPU v6e (Trillium)
En savoir plus sur Cloud TPU

Mettre en service Qwen3-30B-A3B sur GKE avec vLLM TorchTPU multi-hôte et désagrégation de disque persistant sur des accélérateurs TPU v6e Restez organisé à l'aide des collections Enregistrez et classez les contenus selon vos préférences.