הצגת Qwen3-30B-A3B ב-GKE עם vLLM TorchTPU בכמה מארחים ופירוק PD במאיצי TPU v6e

במדריך הזה נסביר איך לפרוס ולהפעיל את המודל Qwen3-30B-A3B (Qwen/Qwen3-30B-A3B-FP8), שהוא מודל שילוב מומחים (MoE), ב-Google Kubernetes Engine ‏(GKE). הפריסה משתמשת במאיצי TPU v6e (Trillium) עם כמה מארחים ובמסגרת להצגת מודלים של vLLM TPU עם TorchTPU (torchtpu) ופירוק של מילוי מראש ופענוח (PD).

בהגדרת הפריסה שמתוארת במדריך הזה, אתם מקצים שני מאגרי צמתים של ct6e-standard-4t עם כמה מארחים, שכל אחד מהם מסודר בטופולוגיה של 2x4 (שמונה שבבי TPU v6e בשני צמתים לכל מאגר צמתים, סה"כ 16 שבבי TPU v6e). אתם משתמשים ב-Kubernetes JobSet API וב-Ray כדי:

  • מריצים שרתי מילוי מראש ופענוח מרובי-מארחים מפורקים בפרוסות TPU v6e נפרדות, כשמקביליות הנתונים ומקביליות המומחים מופעלות.
  • הבקשות מנותבות דרך שרת proxy לביטול צבירת נתונים.

כדי להריץ תהליך עבודה דומה של הצגת מודל vLLM TorchTPU עם כמה מארחים במאיצי TPU7x (Ironwood) עם המודל הגדול יותר Qwen3.5-397B, קראו את המאמר הצגת Qwen3.5-397B ב-GKE עם vLLM TorchTPU עם כמה מארחים במאיצי TPU7x.

מטרות

  • מגדירים את אשכול GKE ואת מאגרי הצמתים של TPU v6e עם כמה מארחים.
  • מעבירים את משקלי המודל לקטגוריה של Cloud Storage.
  • פורסים את שרתי ה-proxy, המילוי מראש והפענוח של vLLM TorchTPU עם מספר מארחים באמצעות JobSet.
  • להריץ עומס עבודה להשוואה של הצגת מודעות ולנתח את התוצאות.
  • מנקים.

עלויות

במסמך הזה משתמשים ברכיבים הבאים שלGoogle Cloud, והשימוש בהם כרוך בתשלום:

כדי להעריך את העלויות בהתאם לשימוש הצפוי, אתם יכולים להשתמש במחשבון העלויות.

כשמסיימים את המשימות שמתוארות במסמך הזה אפשר למחוק את המשאבים שיצרתם כדי להימנע מחיובים נוספים. מידע נוסף זמין בקטע הסרת המשאבים.

לפני שמתחילים

מגדירים את הסביבה למדריך:

  • מפעילים את Kubernetes Engine API בפרויקט ב- Google Cloud :

    הפעלת Kubernetes Engine API

  • מתקינים ומפעילים את Google Cloud CLI, ומתקינים את kubectl. אם אתם משתמשים ב-Cloud Shell, שני הכלים מותקנים מראש.

  • מוודאים שבפרויקט יש מספיק מכסת SSD ומכסת כתובות IP חיצוניות באזור שבו נמצאת קיבולת ה-TPU:

    1. במסוף Google Cloud , נכנסים לדף Quotas & System Limits.

      כניסה לדף Quotas & System Limits

    2. מסננים לפי Persistent Disk SSD (GB)‎ ומוודאים שיש מספיק מכסת SSD זמינה באזור היעד.

    3. מסננים לפי In-use regional external IPv4 addresses ומוודאים שיש לכם מספיק מכסה של כתובות IP באזור היעד.

איך מקבלים טוקן גישה ל-Hugging Face

כדי להשתמש במדריך הזה, אתם צריכים טוקן גישה ל-Hugging Face. אפשר להירשם לחשבון בחינם ב-Hugging Face. אחרי שיוצרים חשבון, יוצרים אסימון גישה:

  1. בדף Welcome to Hugging Face, לוחצים על הדמות של החשבון ובוחרים באפשרות טוקן גישה.
  2. בדף טוקני גישה, לוחצים על Create new token.
  3. בוחרים את סוג הטוקן Read ומזינים שם לטוקן.
  4. טוקן הגישה שלכם יוצג. שומרים את הטוקן במקום בטוח.

התפקידים הנדרשים

כדי לקבל את ההרשאות שדרושות להשלמת המדריך, בקשו מהאדמין להקצות לכם את התפקידים הבאים ב-IAM בפרויקט:

להסבר על מתן תפקידים, ראו איך מנהלים את הגישה ברמת הפרויקט, התיקייה והארגון.

יכול להיות שאפשר לקבל את ההרשאות הנדרשות גם באמצעות תפקידים בהתאמה אישית או תפקידים מוגדרים מראש.

הגדרת משתני הסביבה

כדי להגדיר את משתני הסביבה:

  1. מגדירים את האימות:

    gcloud auth login
gcloud auth application-default login

  2. מייצאים את משתני הסביבה שבהם משתמשים במדריך הזה:

    export PROJECT_ID="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export MACHINE_TYPE="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 .
    • ‫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 שבה אתם משתמשים כדי לאחסן נתונים במחסן ביניים (Stage) את משקלי המודל.
    • ‫YOUR_HF_TOKEN: טוקן הגישה שלכם ל-Hugging Face.

    אם אתם משתמשים בקובץ אימג' של קונטיינר vLLM TorchTPU בהתאמה אישית, החליפו את ה-URI של קובץ האימג' שמוגדר כברירת מחדל (vllm/vllm-torchtpu) במשתנה VLLM_IMAGE ב-URI של קובץ האימג' שלכם. תמונות ציבוריות זמינות במאגר vLLM TorchTPU.

יצירת אשכול GKE

כדי ליצור אשכול GKE במצב Standard ולצרף שני מאגרי צמתים של TPU v6e עם כמה מארחים להפרדה של מילוי מראש ופענוח, מבצעים את השלבים הבאים:

  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. יוצרים שני מאגרי צמתים של TPU v6e עם כמה מארחים, שכל אחד מהם מכיל שני צמתים של ct6e-standard-4t (שמונה שבבי TPU v6e בטופולוגיה של 2x4 לכל מאגר צמתים, ובסך הכול 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‏ (--spot). יכול להיות שייקח זמן ליצור מכונות שמשתמשות בקיבולת של Spot אחרי שיוצרים את מאגר הצמתים, ויכול להיות שהן יופקעו במהלך הרצת עומסי העבודה. אפשר גם ליצור את מאגרי הצמתים באמצעות מקום שמור או Flex-start VM, על ידי העברת הדגלים הבאים:

    • ‫Specific reservation: יצירת מאגר הצמתים באמצעות בקשה ספציפית לשמירת מקום:

      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 VMs: הקצאת מאגר הצמתים באמצעות אפשרות הצריכה 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.

    הפלט של כל מאגר צמתים אמור להיראות כך:

    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. נותנים הרשאות לניהול זהויות והרשאות גישה (IAM) בקטגוריה לחשבון השירות שמוגדר כברירת מחדל ב-Compute Engine ולחשבון השירות שמוגדר כברירת מחדל ב-Kubernetes:

    PROJECT_NUMBER=$(gcloud projects describe "${PROJECT_ID}" --format="value(projectNumber)")

gcloud storage buckets add-iam-policy-binding "gs://${GCS_BUCKET}" \
    --member="serviceAccount:${PROJECT_NUMBER}-compute@developer.gserviceaccount.com" \
    --role="roles/storage.objectAdmin"

gcloud storage buckets add-iam-policy-binding "gs://${GCS_BUCKET}" \
    --member="principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/default/sa/default" \
    --role="roles/storage.admin"

  3. יוצרים סוד של Kubernetes בשם hf-secret כדי לאחסן את טוקן הגישה של Hugging Face:

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

  4. מאחזרים את המניפסט של download_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 לנפח ה-staging ומעבירים אותם לקטגוריה של Cloud Storage:

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

    השלב הזה נמשך כ-5 עד 10 דקות.

  6. מחכים עד שה-Pod יתחיל ועוקבים אחרי יומני ה-Job של ההורדה:

    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) מגדיר משאב JobSet של Kubernetes‏ (vllm-qwen3-30b-pd-disagg-v6e) עם שלוש משימות משוכפלות שמתאמות בין מילוי מראש מבוזר לבין פענוח של הצגת מודעות בשני מאגרי הצמתים של 2x4 TPU v6e:

  • ‫שרת מילוי מראש (prefill): מפעיל שתי רפליקות במאגר צמתים אחד של 2x4 TPU v6e (שני צמתים של ct6e-standard-4t עם ארבעה צ'יפים לכל מארח, סה"כ שמונה ליבות TPU). ה-Pods יוצרים אשכול Ray עם שני צמתים ומפעילים את vLLM TorchTPU ביציאה 8000 עבור Qwen/Qwen3-30B-A3B-FP8. השרת משתמש ב-TPURaidenConnector שמוגדר כ-kv_producer עם הדגלים הבאים:

    • --data-parallel-size=8
    • --tensor-parallel-size=1
    • --enable-expert-parallel
    • --load-format=runai_streamer

  • שרת פענוח (decode): מריץ שתי רפליקות במאגר הצמתים השני של 2x4 TPU v6e (סה"כ שמונה ליבות TPU). ה-JobSet משתמש בכללי קרבה ואי-קרבה של Pod בתווית cloud.google.com/gke-nodepool כדי לבודד את המילוי מראש והפענוח למאגרי צמתים נפרדים. ה-Pods יוצרים אשכול Ray עם שני צמתים ומפעילים את vLLM TorchTPU ביציאה 7200, כאשר TPURaidenConnector מוגדר כ-kv_consumer.

  • שרת proxy להפרדה (proxy): מריץ את הסקריפט toy_proxy_server.py ביציאה 10000 כדי לתאם את ניתוב הבקשות ואת העברת המטמון של KV בין שרת המילוי מראש (יציאה 8000) לבין שרת הפענוח (יציאה 7200).

כדי לפרוס את שרתי ה-vLLM המבוזרים ולעקוב אחריהם:

  1. מתקינים באשכול את ה-CustomResourceDefinition והבקר של Kubernetes JobSet:

    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 (שני צמתים) למילוי מראש ובפרוסה אחת של 2x4 TPU v6e (שני צמתים) לפענוח. כדי לשנות את מספר הצמתים בכל אחד מהמאגרים, מעדכנים את השדה replicas (replicas: NUM_NODES). מחליפים את NUM_NODES במספר הצמתים שרוצים להשתמש בהם.

  4. עוקבים אחרי יומני ה-Pod של המילוי מראש והפענוח עד שהאפליקציה מסיימת את ההפעלה:

    # Prefill logs
kubectl logs -l jobset.sigs.k8s.io/replicatedjob-name=prefill -f --tail=-1

# Decode logs
kubectl logs -l jobset.sigs.k8s.io/replicatedjob-name=decode -f --tail=-1

    הטעינה וההידור הראשוניים של המודל נמשכים כ-15 עד 25 דקות בהפעלה הראשונה.

    אם אתם משתמשים בקיבולת שניתנת להפסקה לפני הזמן, כמו קיבולת של Spot או מכונות וירטואליות מסוג Flex-start, יכול להיות שייקח יותר זמן ליצור את הצמתים של פרוסות ה-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.

הרצת עומס עבודה להערכת ביצועים

כדי להשוות בין מודל Qwen3-30B-A3B שנפרס באמצעות עומס עבודה של קלט של 8,192 טוקנים ופלט של 1,024 טוקנים (8k/1k) לבין שרת ה-proxy להפרדה, מבצעים את השלבים הבאים:

  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. מחכים שה-Pod של נקודת ההשוואה יהיה מוכן, ועוקבים אחרי היומנים של משימת נקודת ההשוואה:

    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

המאמרים הבאים