Google verwendet KI-Technologie, um Inhalte in Ihre bevorzugte Sprache zu übersetzen. KI-Übersetzungen können Fehler enthalten.

LLMs auf TPUs mit mehreren Hosts mit Ray und TorchTPU bereitstellen

Autopilot Standard

In dieser Anleitung erfahren Sie, wie Sie ein Mixture-of-Experts-Modell (MoE) vom Typ Gemma 4 26B-A4B-it (Large Language Model, LLM) mit Tensor- und Expertenparallelität auf Multi-Host-Cloud TPU-Slices in Google Kubernetes Engine (GKE) bereitstellen. Dazu verwenden Sie Ray Serve LLM, vLLM und vLLM-TorchTPU.

Diese Anleitung richtet sich an Machine-Learning-Entwickler, Plattformadministratoren und -operatoren sowie Daten- und KI-Spezialisten, die die Kubernetes-Containerorchestrierung verwenden, um KI- und ML-Arbeitslasten auf verteilten TPU-Slices mit mehreren Hosts bereitzustellen. Weitere Informationen zu gängigen Rollen und Beispielaufgaben, auf die wir in Google Cloud -Inhalten verweisen, finden Sie unter Häufig verwendete GKE-Nutzerrollen und -Aufgaben.

Bevor Sie mit dieser Anleitung beginnen, sollten Sie mit Folgendem vertraut sein:

Hintergrund

In diesem Abschnitt werden die in dieser Anleitung verwendeten Schlüsseltechnologien beschrieben.

TPUs

Tensor Processing Units (TPUs) sind von Google speziell entwickelte Beschleuniger für Arbeitslasten des maschinellen Lernens. Diese Anleitung unterstützt sowohl TPU Trillium (v6e) als auch Ironwood (TPU7x). In einem TPU-Slice mit mehreren Hosts kommunizieren mehrere TPU-Hosts direkt über ein ICI-Netzwerk (Inter-Chip Interconnect) mit hoher Geschwindigkeit. Dadurch werden der aggregierte High Bandwidth Memory (HBM) und die kollektive Kommunikation mit niedriger Latenz bereitgestellt, die für die Bereitstellung von Modellen erforderlich sind, die die Speicherkapazität eines einzelnen Chips oder TPU-Hosts überschreiten.

vLLM und TorchTPU

vLLM ist eine Open-Source-Inferenz- und Bereitstellungs-Engine für LLMs. TorchTPU ist ein Open-Source-PyTorch-Backend für Cloud TPU, das Geräteverwaltung und kollektive Kommunikation über ICI bietet.

vLLM-TorchTPU registriert das TPU-Plattform-Plug-in bei vLLM und kompiliert PyTorch-Modell-Forward-Passes für Cloud TPU über TorchTPU und den XLA-Compiler (Accelerated Linear Algebra). Es unterstützt verteilte Tensor-Parallelität und Experten-Parallelität über TPU-Slices mit mehreren Hosts hinweg mithilfe von Ray.

Ray Serve-LLM in GKE

Ray Serve LLM bietet einen OpenAI-kompatiblen HTTP-Ingress für /v1/chat/completions-, /v1/completions- und /v1/models-Endpunkte und orchestriert verteilte vLLM-Bereitstellungen in der GKE. Wenn Sie das Ray on GKE-Add-on aktivieren, fügen der KubeRay-Operator und sein TPU-Webhook Slice-Topologiemetadaten, Worker-Hostnamen und TorchTPU-Koordinationsvariablen in jeden TPU-Worker-Pod mit mehreren Hosts ein.

Ray Serve LLM bietet die folgenden Funktionen für die Bereitstellung von TPUs mit mehreren Hosts:

Bei der topologiebasierten Slice-Platzierung wird für jedes Modellreplikat ein intaktes Multi-Host-TPU-Slice reserviert, wenn Sie im Feld accelerator_config eine TPU-Topologie angeben. Dabei wird ein vLLM-Worker pro TPU-Gerät im Slice geplant.
Präfixbasiertes Routing leitet Anfragen mit freigegebenen System-Prompts oder Unterhaltungsverläufen mit mehreren Runden an das Replikat weiter, das diese Präfix-Schlüssel-Wert-Blöcke (KV) bereits im TPU-HBM zwischenspeichert, wenn die Bereitstellung über mehrere TPU-Slice-Replikate erfolgt.
Beim direkten Token-Streaming werden ausgehende Token-Streams über gleichzeitige Anfragen hinweg zusammengefasst und Antworten direkt von Modellreplikaten gestreamt, um die Zeit bis zum ersten Token (Time-To-First-Token, TTFT) und die Latenz zwischen Tokens (Inter-Token Latency, ITL) zu reduzieren.
Beim gleichzeitigkeitsbasierten Autoscaling von Replikaten werden Modellreplikate basierend auf der aktuellen Anfragegleichzeitigkeit in ganzen Slice-Schritten skaliert.

Ziele

In dieser Anleitung wird Folgendes beschrieben:

Erstellen Sie ein Container-Image, das Ray Serve über das vordefinierte vLLM-TorchTPU-Basis-Image legt, und übertragen Sie es per Push an Artifact Registry.
Erstellen Sie einen GKE-Cluster im Autopilot- oder Standardmodus auf TPU Trillium (v6e) oder Ironwood (TPU7x).
Stellen Sie eine benutzerdefinierte RayCluster-Ressource mit mehreren Hosts bereit und prüfen Sie die TPU-Topologiekonfiguration.
Stellen Sie eine Ray Serve-LLM-Anwendung bereit, um das google/gemma-4-26B-A4B-it-Modell mit Tensor-Parallelität und Experten-Parallelität über ein TPU-Slice mit mehreren Hosts bereitzustellen, oder stellen Sie sowohl den Cluster als auch die Anwendung deklarativ mit einer benutzerdefinierten RayService-Ressource bereit.
Senden Sie OpenAI-kompatible Anfragen zur Modellermittlung und Chatvervollständigung an das bereitgestellte Modell.

Kosten

In dieser Anleitung verwenden Sie die folgenden kostenpflichtigen Komponenten von Google Cloud:

Mit dem Preisrechner können Sie eine Kostenschätzung für Ihre voraussichtliche Nutzung vornehmen.

Nach Abschluss der in dieser Anleitung beschriebenen Aufgaben können Sie weitere Kosten vermeiden, indem Sie die erstellten Ressourcen löschen. Weitere Informationen finden Sie unter Bereinigen.

Hinweis

Führen Sie die folgenden Aufgaben aus, bevor Sie beginnen:

Aktivieren Sie die Artifact Registry API und die Google Kubernetes Engine API.

APIs aktivieren

Wenn Sie die Google Cloud CLI für diese Aufgabe verwenden möchten, müssen Sie die gcloud CLI installieren und dann initialisieren. Wenn Sie die gcloud CLI bereits installiert haben, rufen Sie die neueste Version mit dem Befehl gcloud components update ab. Ältere gcloud CLI-Versionen unterstützen möglicherweise nicht die Ausführung der Befehle in diesem Dokument.
Hinweis: Legen Sie für vorhandene Installationen der gcloud CLI das compute/region-Attribut fest. Wenn Sie hauptsächlich zonale Cluster verwenden, legen Sie stattdessen compute/zone fest. Wenn Sie einen standardmäßigen Speicherort festlegen, können Sie in der gcloud CLI Fehler wie diesen vermeiden: One of [--zone, --region] must be supplied: Please specify location. Bei bestimmten Befehlen müssen Sie möglicherweise den Speicherort angeben, wenn er sich von dem von Ihnen festgelegten Standard unterscheidet.

Prüfen Sie, ob Ihr Projekt über ein ausreichendes Kontingent für TPU Trillium (v6e)- oder Ironwood (TPU7x)-Kapazität in der ausgewählten Zone verfügt. Weitere Informationen finden Sie unter Cloud TPU-Kontingente.
Sie benötigen ein Hugging Face-Konto. Akzeptieren Sie die Lizenzbedingungen auf der google/gemma-4-26B-A4B-it-Modellseite und generieren Sie ein Hugging Face-Zugriffstoken mit Read-Berechtigungen.
Prüfen Sie, ob Ihr Konto die folgenden IAM-Rollen hat:
- Kubernetes Engine-Administrator (roles/container.admin)
- Artifact Registry-Administrator (roles/artifactregistry.admin)

Umgebung vorbereiten

In dieser Anleitung verwenden Sie Cloud Shell oder eine lokale Workstation, um Ressourcen zu verwalten, die in Google Cloudgehostet werden. Die in dieser Anleitung verwendeten Befehlszeilentools sind in Cloud Shell vorinstalliert, einschließlich kubectl, docker, python3 und der gcloud-Befehlszeile.

So richten Sie Ihre Umgebung mit Cloud Shell ein:

Starten Sie in der Google Cloud -Console eine Cloud Shell-Sitzung. Klicken Sie dazu in der Google Cloud -Console auf Cloud Shell aktivieren. Dadurch wird im unteren Bereich der Google Cloud Console eine Sitzung gestartet.

Klonen Sie das Repository kubernetes-engine-samples und wechseln Sie in das Beispielverzeichnis für diese Anleitung:

git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples.git
cd kubernetes-engine-samples/ai-ml/gke-ray/rayserve/llm/torchtpu

Erstellen und aktivieren Sie eine lokale virtuelle Python-Umgebung mit installierter Ray-Befehlszeile:
```
python3 -m venv ~/ray-env
source ~/ray-env/bin/activate
pip install "ray[default]==2.58.0"
```
Legen Sie die Umgebungsvariablen für Ihr Projekt und Ihren Cluster in Google Cloud fest:
```
export PROJECT_ID=PROJECT_ID
export CLUSTER_NAME=CLUSTER_NAME
export REGION=REGION
export ZONE=ZONE
export ARTIFACT_REGISTRY=ARTIFACT_REGISTRY
export HF_TOKEN=HF_TOKEN
export IMAGE_URI=${REGION}-docker.pkg.dev/${PROJECT_ID}/${ARTIFACT_REGISTRY}/ray-torchtpu-vllm:latest
gcloud config set project ${PROJECT_ID}
```
Ersetzen Sie Folgendes:
- PROJECT_ID: Projekt-ID in Google Cloud .
- CLUSTER_NAME: der Name Ihres GKE-Cluster, z. B. ray-serve-torchtpu-cluster.
- REGION: die Region, in der sich Ihr Artifact Registry-Repository und Ihre TPU-Kapazität befinden, z. B. us-east5.
- ZONE: Die Zone, in der Ihre TPU-Kapazität verfügbar ist, z. B. us-east5-b. Weitere Informationen finden Sie unter TPU-Verfügbarkeit in GKE prüfen.
- ARTIFACT_REGISTRY: der Name Ihres Artifact Registry-Docker-Repositorys, z. B. ray-torchtpu.
- HF_TOKEN: Ihr Hugging Face-Zugriffstoken mit Read-Berechtigungen.

Benutzerdefiniertes Container-Image erstellen

Erstellen Sie ein Container-Image, das das Paket ray[default,serve] zusätzlich zum vordefinierten vLLM-TorchTPU-Basis-Image installiert, das TorchTPU, vLLM und das vLLM-TorchTPU-Plug-in enthält.

Prüfen Sie die Dockerfile im Verzeichnis kubernetes-engine-samples/ai-ml/gke-ray/rayserve/llm/torchtpu:

FROM vllm/vllm-torchtpu:v0.29.0

ENV HF_HOME=/dev/shm/huggingface
ENV RAY_SERVE_DEFAULT_HTTP_HOST=0.0.0.0

# Install wget for KubeRay health probes and pin Ray Serve and transformers
# versions compatible with vLLM 0.29.0 and multi-host TPUs.
RUN apt-get update && apt-get install -y --no-install-recommends wget \
    && rm -rf /var/lib/apt/lists/* \
    && pip install --no-cache-dir \
       "ray[default,serve] @ https://s3-us-west-2.amazonaws.com/ray-wheels/master/6b2b414bb7a5b4e2705ed9f4fdab9787b92f3b08/ray-3.0.0.dev0-cp312-cp312-manylinux2014_x86_64.whl" \
       "transformers>=5.18.0"

Für die TPU-Bereitstellung mit mehreren Hosts mit Ray Serve LLM ist Ray Version 2.56.0 oder höher auf TPU Trillium (v6e) und Ray Version 2.58.0 oder höher für die Planung pro Chiplet auf Ironwood (TPU7x) erforderlich.

Erstellen Sie das Docker-Repository in Artifact Registry, erstellen Sie das Container-Image und übertragen Sie es per Push an Artifact Registry:

gcloud artifacts repositories create ${ARTIFACT_REGISTRY} \
    --repository-format=docker \
    --location=${REGION} \
    --description="Ray and vLLM-TorchTPU container images"

gcloud auth configure-docker ${REGION}-docker.pkg.dev,us-docker.pkg.dev
docker build -t ${IMAGE_URI} .
docker push ${IMAGE_URI}

GKE-Cluster erstellen

Sie können LLMs auf TPU-Slices mit mehreren Hosts in einem GKE Autopilot- oder Standardcluster mit aktiviertem verwalteten Ray-Operator-Add-on bereitstellen.

Erstellen Sie Ihren GKE-Cluster und konfigurieren Sie die Anmeldedaten für das kubectl-Befehlszeilentool:
Autopilot
1. Erstellen Sie einen regionalen GKE Autopilot-Cluster mit aktiviertem Ray-Operator-Add-on:
  
  gcloud container clusters create-auto ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --region=${REGION} \ --cluster-version=1.36 \ --enable-ray-operator
  
  Wenn Sie die benutzerdefinierte Ressource „RayCluster“ oder „RayService“ in einem späteren Schritt bereitstellen, stellt GKE Autopilot automatisch den CPU-Head-Knoten und das TPU-Slice mit mehreren Hosts bereit, die Ihren Pod-Ressourcenanfragen und Knotenauswahlen entsprechen.
2. Konfigurieren Sie das kubectl-Befehlszeilentool für die Kommunikation mit Ihrem Cluster:
  
  gcloud container clusters get-credentials ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --region=${REGION}
Standard
1. Erstellen Sie einen zonalen GKE Standard-Cluster mit einem CPU-Knotenpool zum Hosten des Ray-Head-Pods, Image-Streaming und aktiviertem GKE Dataplane V2:
  
  gcloud container clusters create ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --cluster-version=1.36 \ --machine-type=n2-standard-8 \ --num-nodes=1 \ --addons=RayOperator \ --enable-image-streaming \ --enable-dataplane-v2
  Hinweis: Wenn Sie Open-Source-KubeRay selbst verwalten, anstatt das GKE-RayOperator-Add-on zu verwenden, installieren Sie KubeRay und den KubeRay-TPU-Webhook in Version 1.4.0 oder höher, um die TorchTPU-Umgebungsvariable einzufügen.
2. Konfigurieren Sie das kubectl-Befehlszeilentool für die Kommunikation mit Ihrem Cluster:
  
  gcloud container clusters get-credentials ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE}
3. Erstellen Sie einen TPU-Slice-Knotenpool mit mehreren Hosts und 16 Geräten mit einer 4x4-Topologie auf vier Hosts auf TPU Trillium (v6e) oder einer 2x2x2-Topologie auf zwei Hosts auf Ironwood (TPU7x):
  
  Trillium (v6e)
  gcloud container node-pools create tpu-multi-host-pool \ --project=${PROJECT_ID} \ --cluster=${CLUSTER_NAME} \ --zone=${ZONE} \ --machine-type=ct6e-standard-4t \ --tpu-topology=4x4 \ --num-nodes=4
  
  Ironwood (TPU7x)
  gcloud container node-pools create tpu-multi-host-pool \ --project=${PROJECT_ID} \ --cluster=${CLUSTER_NAME} \ --zone=${ZONE} \ --machine-type=tpu7x-standard-4t \ --tpu-topology=2x2x2 \ --num-nodes=2
  
  Tipp: Wenn Sie eine bestimmte TPU-Reservierung nutzen möchten, hängen Sie die Flags --reservation-affinity=specific und --reservation=RESERVATION_NAME an den Befehl node-pools create an, wobei RESERVATION_NAME der Name Ihrer TPU-Reservierung ist. Wenn Sie Spot-VMs verwenden möchten, hängen Sie das Flag --spot an.
Erstellen Sie ein Kubernetes-Secret mit Ihrem Hugging Face-Token, damit sich Ray-Pods beim Herunterladen von Modellgewichtungen bei Hugging Face authentifizieren können:
```
kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN}
```

Benutzerdefinierte RayCluster-Ressource mit mehreren Hosts bereitstellen

Stellen Sie eine benutzerdefinierte RayCluster-Ressource bereit, die einen CPU-Head-Pod, ein In-Memory-/dev/shm-Volume für den Ray-Objektspeicher mit gemeinsamem Arbeitsspeicher und eine Multi-Host-TPU-Worker-Gruppe bereitstellt.

Sehen Sie sich das RayCluster-Manifest für Ihre TPU-Version im Beispiel-Repository an:

Trillium (v6e)

apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ray-serve-torchtpu
spec:
  headGroupSpec:
    rayStartParams:
      num-cpus: "0"
    template:
      spec:
        containers:
        - name: ray-head
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          ports:
          - containerPort: 6379
            name: gcs
          - containerPort: 8265
            name: dashboard
          - containerPort: 10001
            name: client
          - containerPort: 8000
            name: serve
          resources:
            limits:
              cpu: "4"
              memory: "16Gi"
            requests:
              cpu: "4"
              memory: "16Gi"
  workerGroupSpecs:
  - groupName: v6e-16
    replicas: 1
    numOfHosts: 4
    template:
      spec:
        containers:
        - name: ray-worker
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          resources:
            limits:
              cpu: "64"
              memory: "200Gi"
              google.com/tpu: "4"
            requests:
              cpu: "64"
              memory: "200Gi"
              google.com/tpu: "4"
          volumeMounts:
          - mountPath: /dev/shm
            name: dshm
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 100Gi
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: 4x4

Ironwood (TPU7x)

apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ray-serve-torchtpu
spec:
  headGroupSpec:
    rayStartParams:
      num-cpus: "0"
    template:
      spec:
        containers:
        - name: ray-head
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          - name: RAY_TPU_RESOURCE_PER_CHIP
            value: "2"
          ports:
          - containerPort: 6379
            name: gcs
          - containerPort: 8265
            name: dashboard
          - containerPort: 10001
            name: client
          - containerPort: 8000
            name: serve
          resources:
            limits:
              cpu: "4"
              memory: "16Gi"
            requests:
              cpu: "4"
              memory: "16Gi"
  workerGroupSpecs:
  - groupName: tpu7x-16
    replicas: 1
    numOfHosts: 2
    resources:
      # Advertises 8 logical TPU cores (2 chiplets per physical chip) per host.
      TPU: "8"
    template:
      spec:
        containers:
        - name: ray-worker
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          - name: RAY_TPU_RESOURCE_PER_CHIP
            value: "2"
          resources:
            limits:
              cpu: "64"
              memory: "250Gi"
              google.com/tpu: "4"
            requests:
              cpu: "64"
              memory: "250Gi"
              google.com/tpu: "4"
          volumeMounts:
          - mountPath: /dev/shm
            name: dshm
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 128Gi
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu7x
          cloud.google.com/gke-tpu-topology: 2x2x2

Auf Ironwood (TPU7x) enthält jeder physische Chip zwei Chiplets, was acht logische TPU-Kerne pro Host mit vier Chips ergibt. Im Manifest wird der Wert des Felds TPU im Abschnitt resources auf "8" festgelegt. Außerdem wird die Umgebungsvariable RAY_TPU_RESOURCE_PER_CHIP auf "2" gesetzt, um acht logische TPU-Ressourcen pro Host anzukündigen und zu planen.

Wenden Sie das Manifest an, um die benutzerdefinierte RayCluster-Ressource zu erstellen:

Trillium (v6e)

envsubst < ray-cluster.tpu-v6e-multihost.yaml | kubectl apply -f -

Ironwood (TPU7x)

envsubst < ray-cluster.tpu-v7x-multihost.yaml | kubectl apply -f -

Warten Sie, bis der Head-Pod und alle TPU-Worker-Pods den Status Running und 1/1 READY erreicht haben:

kubectl get pods -l ray.io/cluster=ray-serve-torchtpu

Die Ausgabe sieht etwa so aus:

Trillium (v6e)

NAME                                     READY   STATUS    RESTARTS   AGE
ray-serve-torchtpu-head-5cj7g            1/1     Running   0          35s
ray-serve-torchtpu-v6e-16-worker-6x8cg   1/1     Running   0          35s
ray-serve-torchtpu-v6e-16-worker-79j5l   1/1     Running   0          35s
ray-serve-torchtpu-v6e-16-worker-psxbl   1/1     Running   0          35s
ray-serve-torchtpu-v6e-16-worker-rz7dv   1/1     Running   0          35s

Ironwood (TPU7x)

NAME                                       READY   STATUS    RESTARTS   AGE
ray-serve-torchtpu-head-w7b2v              1/1     Running   0          2m30s
ray-serve-torchtpu-tpu7x-16-worker-6zlnj   1/1     Running   0          2m30s
ray-serve-torchtpu-tpu7x-16-worker-cb4zz   1/1     Running   0          2m30s

Prüfen Sie, ob der KubeRay-TPU-Webhook die Umgebungsvariablen für das Multi-Host-Slice, die Worker-Koordination und TorchTPU in die Worker-Pods eingefügt hat:

kubectl get pods -l ray.io/node-type=worker \
    -o=jsonpath='{range .items[0].spec.containers[0].env[*]}{.name}={.value}{"\n"}{end}' \
    | grep -E "TPU_WORKER_ID|TPU_WORKER_HOSTNAMES|TPU_NAME|TORCH_TPU_TOPOLOGY|TORCH_TPU_SLICEBUILDER_ADDRESSES"

Die Ausgabe sieht etwa so aus:

Trillium (v6e)

TPU_WORKER_HOSTNAMES=v6e-16-0-0.ray-serve-torchtpu-headless,v6e-16-0-1.ray-serve-torchtpu-headless,v6e-16-0-2.ray-serve-torchtpu-headless,v6e-16-0-3.ray-serve-torchtpu-headless
TORCH_TPU_TOPOLOGY=4,4,1
TORCH_TPU_SLICEBUILDER_ADDRESSES=v6e-16-0-0.ray-serve-torchtpu-headless:8471,v6e-16-0-0.ray-serve-torchtpu-headless:8472,v6e-16-0-0.ray-serve-torchtpu-headless:8473,v6e-16-0-0.ray-serve-torchtpu-headless:8474,v6e-16-0-1.ray-serve-torchtpu-headless:8471,v6e-16-0-1.ray-serve-torchtpu-headless:8472,v6e-16-0-1.ray-serve-torchtpu-headless:8473,v6e-16-0-1.ray-serve-torchtpu-headless:8474,v6e-16-0-2.ray-serve-torchtpu-headless:8471,v6e-16-0-2.ray-serve-torchtpu-headless:8472,v6e-16-0-2.ray-serve-torchtpu-headless:8473,v6e-16-0-2.ray-serve-torchtpu-headless:8474,v6e-16-0-3.ray-serve-torchtpu-headless:8471,v6e-16-0-3.ray-serve-torchtpu-headless:8472,v6e-16-0-3.ray-serve-torchtpu-headless:8473,v6e-16-0-3.ray-serve-torchtpu-headless:8474
TPU_WORKER_ID=1
TPU_NAME=v6e-16-0

Ironwood (TPU7x)

TPU_WORKER_HOSTNAMES=tpu7x-16-0-0.ray-serve-torchtpu-headless,tpu7x-16-0-1.ray-serve-torchtpu-headless
TORCH_TPU_TOPOLOGY=2,2,2,2
TORCH_TPU_SLICEBUILDER_ADDRESSES=tpu7x-16-0-0.ray-serve-torchtpu-headless:8471,tpu7x-16-0-0.ray-serve-torchtpu-headless:8472,tpu7x-16-0-0.ray-serve-torchtpu-headless:8473,tpu7x-16-0-0.ray-serve-torchtpu-headless:8474,tpu7x-16-0-0.ray-serve-torchtpu-headless:8475,tpu7x-16-0-0.ray-serve-torchtpu-headless:8476,tpu7x-16-0-0.ray-serve-torchtpu-headless:8477,tpu7x-16-0-0.ray-serve-torchtpu-headless:8478,tpu7x-16-0-1.ray-serve-torchtpu-headless:8471,tpu7x-16-0-1.ray-serve-torchtpu-headless:8472,tpu7x-16-0-1.ray-serve-torchtpu-headless:8473,tpu7x-16-0-1.ray-serve-torchtpu-headless:8474,tpu7x-16-0-1.ray-serve-torchtpu-headless:8475,tpu7x-16-0-1.ray-serve-torchtpu-headless:8476,tpu7x-16-0-1.ray-serve-torchtpu-headless:8477,tpu7x-16-0-1.ray-serve-torchtpu-headless:8478
TPU_WORKER_ID=1
TPU_NAME=tpu7x-16-0

Sie können die TPU-Worker-Knoten und ihre automatisch erkannten ray.io/tpu-*-Labels auf dem Tab Cluster des Ray-Dashboards prüfen.

Leiten Sie lokale Ports an das Ray-Dashboard auf Port 8265 und den Ray Serve-HTTP-Endpunkt auf Port 8000 weiter:
```
kubectl port-forward svc/ray-serve-torchtpu-head-svc 8265:8265 8000:8000 >/dev/null 2>&1 &
```

Ray Serve-LLM-Anwendung erstellen und bereitstellen

Stellen Sie die Ray Serve-LLM-Anwendung in Ihrer benutzerdefinierten RayCluster-Ressource bereit, um das google/gemma-4-26B-A4B-it-Modell über das TPU-Slice mit mehreren Hosts bereitzustellen.

Sehen Sie sich das Script serve_vllm_torchtpu.py im Beispiel-Repository an:

from ray.serve.deployment import Application
from ray.serve.llm import (
    LLMConfig,
    LLMServingArgs,
    ModelLoadingConfig,
    build_openai_app,
)


def app(args: dict[str, str]) -> Application:
    llm_config = LLMConfig(
        model_loading_config=ModelLoadingConfig(
            model_id=args.get("model_id", "google/gemma-4-26B-A4B-it"),
        ),
        accelerator_type=args.get("accelerator_type", "TPU-V6E"),
        accelerator_config={
            "kind": "tpu",
            "topology": args.get("topology", "4x4"),
        },
        placement_group_config={
            "bundle_per_worker": {"TPU": 1},
        },
        engine_kwargs={
            "tensor_parallel_size": int(args.get("tensor_parallel_size", "16")),
            "enable_expert_parallel": True,
            "enable_prefix_caching": True,
            "language_model_only": True,
            "disable_hybrid_kv_cache_manager": True,
            "max_model_len": int(args.get("max_model_len", "8192")),
            "max_num_batched_tokens": int(
                args.get("max_num_batched_tokens", "4096")
            ),
        },
        experimental_configs={
            "stream_batching_interval_ms": 10,
        },
        runtime_env={
            "env_vars": {
                "TPU_MULTIHOST_BACKEND": "ray",
                # Execute Gemma 4 routed expert layers on TPU TensorCores instead of SparseCore.
                "USE_MOE_SPARSE_CORE": "0",
                "VLLM_RAY_EXTRA_ENV_VARS_TO_COPY": "USE_MOE_SPARSE_CORE",
            }
        },
    )
    return build_openai_app(LLMServingArgs(llm_configs=[llm_config]))

Das Skript konfiguriert das LLMConfig-Objekt mit den folgenden Einstellungen für die Multi-Host-TPU-Bereitstellung:

Mit den Feldern accelerator_config und placement_group_config wird eine Platzierungsgruppe auf einem einzelnen TPU-Slice reserviert und ein TPU-Kern pro vLLM-Worker auf dem Slice zugewiesen.
Das Feld engine_kwargs fragmentiert das Modell mithilfe von Tensor- und Expertenparallelität über alle 16 TPU-Kerne und aktiviert das Präfix-Caching. Außerdem werden multimodale Eingaben deaktiviert, um TPU-HBM für die Textinferenz zu reservieren. Es wird ein einheitliches KV-Cache-Layout für Gemma 4-Attention-Layer verwendet und die Kontextlänge auf 8.192 Tokens und die Batchgröße für die Vorbefüllung auf 4.096 Tokens begrenzt.
Im Feld experimental_configs werden ausgehende Tokenstreams alle 10 ms zusammengefasst, um den Streaming-Overhead bei gleichzeitigen Anfragen zu reduzieren.
Das Feld runtime_env aktiviert das Ray-Multi-Host-Backend für vLLM-TorchTPU (TPU_MULTIHOST_BACKEND: "ray"), führt Gemma 4 MoE-Expertenschichten auf TPU-TensorCores anstelle von SparseCore aus und gibt diese Einstellung an alle 16 Worker-Akteure weiter.

Stellen Sie die Ray Serve-Anwendung mit dem Befehl serve run in der benutzerdefinierten RayCluster-Ressource bereit. Der Befehl streamt die Bereitstellungslogs, bis die Modellgewichtungen vollständig geladen sind, die XLA-Grafikkompilierung auf allen 16 TPU-Geräten abgeschlossen ist und die Serve-Anwendung den Status RUNNING erreicht:

Trillium (v6e)

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- serve run --non-blocking serve_vllm_torchtpu:app accelerator_type=TPU-V6E topology=4x4

Ironwood (TPU7x)

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- serve run --non-blocking serve_vllm_torchtpu:app accelerator_type=TPU-V7X topology=2x2x2

Die Ausgabe sieht etwa so aus:

Trillium (v6e)

(ServeController pid=1637) INFO 2026-10-13 10:23:05,160 controller 1637 -- Deploying new version of Deployment(name='LLMServer:google--gemma-4-26B-A4B-it', app='default') (initial target replicas: 1).
(ServeController pid=1637) INFO 2026-10-13 10:23:05,161 controller 1637 -- Deploying new version of Deployment(name='OpenAiIngress', app='default') (initial target replicas: 1).
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) INFO 10-13 10:24:06 [__init__.py:318] TPU info: node_name=v6e-16-0 | tpu_type=v6e-16 | worker_id=1 | num_chips=4 | num_cores_per_chip=1
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) INFO 10-13 10:24:06 [__init__.py:238] Platform plugin torchtpu is activated
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) INFO 10-13 10:24:08 [core.py:112] Initializing a V1 LLM engine (v0.22.1) with config: model='google/gemma-4-26B-A4B-it', speculative_config=None, tokenizer='google/gemma-4-26B-A4B-it', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.bfloat16, max_seq_len=8192, download_dir=None, load_format=auto, tensor_parallel_size=16, pipeline_parallel_size=1, data_parallel_size=1, decode_context_parallel_size=1, dcp_comm_backend=ag_rs, disable_custom_all_reduce=True, quantization=None, quantization_config=None, enforce_eager=False, enable_return_routed_experts=False, kv_cache_dtype=auto, device_config=tpu, ...
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) (RayWorkerWrapper pid=1017, ip=10.72.14.21) INFO 10-13 10:24:30 [expert_map_manager.py:245] [EP Rank 4/16] Expert parallelism is enabled. Expert placement strategy: linear. Local/global number of experts: 8/128. Experts local to global index map: 0->32, 1->33, 2->34, 3->35, 4->36, 5->37, 6->38, 7->39.
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) (RayWorkerWrapper pid=856, ip=10.72.6.21) INFO 10-13 10:27:55 [default_loader.py:397] Loading weights took 12.76 seconds
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) (RayWorkerWrapper pid=856, ip=10.72.6.21) INFO 10-13 10:32:53 [tpu_worker.py:376] Memory statistics | total_hbm_limit_gb=31.25GiB | total_hbm_limit_cap_gb=28.75GiB | total_hbm_used_gb=3.69GiB | kv_cache_headroom_gb=5.0GiB | kv_connector_hbm_reserve_gb=0.0GiB | total_hbm_avail_gb=20.06GiB
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) INFO 10-13 10:32:54 [kv_cache_utils.py:1733] GPU KV cache size: 600,832 tokens
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) INFO 10-13 10:32:54 [kv_cache_utils.py:1734] Maximum concurrency for 8,192 tokens per request: 73.34x
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) INFO 2026-10-13 10:36:59,851 default_LLMServer:google--gemma-4-26B-A4B-it b2grwv9c -- Started vLLM engine.
INFO 2026-10-13 10:37:05,997 serve 1456 -- Application 'default' is ready at http://127.0.0.1:8000/.
2026-10-13 10:37:11,447 SUCC cli.py:66 -- ------------------------------------------
2026-10-13 10:37:11,447 SUCC cli.py:67 -- Job 'raysubmit_94DTwVzR2pvGQVwR' succeeded
2026-10-13 10:37:11,447 SUCC cli.py:68 -- ------------------------------------------

Ironwood (TPU7x)

(ServeController pid=2042) INFO 2026-10-13 10:24:38,092 controller 2042 -- Deploying new version of Deployment(name='LLMServer:google--gemma-4-26B-A4B-it', app='default') (initial target replicas: 1).
(ServeController pid=2042) INFO 2026-10-13 10:24:38,093 controller 2042 -- Deploying new version of Deployment(name='OpenAiIngress', app='default') (initial target replicas: 1).
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) INFO 10-13 10:24:49 [__init__.py:318] TPU info: node_name=tpu7x-16-0 | tpu_type=tpu7x-16 | worker_id=0 | num_chips=8 | num_cores_per_chip=2
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) INFO 10-13 10:24:49 [__init__.py:238] Platform plugin torchtpu is activated
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) INFO 10-13 10:25:23 [core.py:112] Initializing a V1 LLM engine (v0.22.1) with config: model='google/gemma-4-26B-A4B-it', speculative_config=None, tokenizer='google/gemma-4-26B-A4B-it', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.bfloat16, max_seq_len=8192, download_dir=None, load_format=auto, tensor_parallel_size=16, pipeline_parallel_size=1, data_parallel_size=1, decode_context_parallel_size=1, dcp_comm_backend=ag_rs, disable_custom_all_reduce=True, quantization=None, quantization_config=None, enforce_eager=False, enable_return_routed_experts=False, kv_cache_dtype=auto, device_config=tpu, ...
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) (RayWorkerWrapper pid=602, ip=172.16.10.6) INFO 10-13 10:25:44 [expert_map_manager.py:245] [EP Rank 8/16] Expert parallelism is enabled. Expert placement strategy: linear. Local/global number of experts: 8/128. Experts local to global index map: 0->64, 1->65, 2->66, 3->67, 4->68, 5->69, 6->70, 7->71.
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) (RayWorkerWrapper pid=762) INFO 10-13 10:29:06 [default_loader.py:397] Loading weights took 11.69 seconds
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) (RayWorkerWrapper pid=604, ip=172.16.10.6) INFO 10-13 10:34:51 [tpu_worker.py:376] Memory statistics | total_hbm_limit_gb=94.75GiB | total_hbm_limit_cap_gb=87.17GiB | total_hbm_used_gb=3.74GiB | kv_cache_headroom_gb=5.0GiB | kv_connector_hbm_reserve_gb=0.0GiB | total_hbm_avail_gb=78.43GiB
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) INFO 10-13 10:34:58 [kv_cache_utils.py:1733] GPU KV cache size: 2,349,568 tokens
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) INFO 10-13 10:34:58 [kv_cache_utils.py:1734] Maximum concurrency for 8,192 tokens per request: 286.81x
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) INFO 2026-10-13 10:40:13,586 default_LLMServer:google--gemma-4-26B-A4B-it a15k5cv3 -- Started vLLM engine.
INFO 2026-10-13 10:40:18,415 serve 1953 -- Application 'default' is ready at http://127.0.0.1:8000/.
2026-10-13 10:40:24,372 SUCC cli.py:66 -- ------------------------------------------
2026-10-13 10:40:24,372 SUCC cli.py:67 -- Job 'raysubmit_mzWQDQGDM9QSgcMN' succeeded
2026-10-13 10:40:24,372 SUCC cli.py:68 -- ------------------------------------------

Bereitgestelltes Modell abfragen

Nachdem die Ray Serve-Anwendung den Status RUNNING erreicht hat, prüfen Sie den registrierten Modellendpunkt, senden Sie eine OpenAI-kompatible /v1/chat/completions-Anfrage über Port 8000 und führen Sie optional einen Benchmark für Durchsatz und Latenz durch.

Prüfen Sie, ob das Modell auf dem Endpunkt /v1/models registriert ist:

curl http://localhost:8000/v1/models

Die Ausgabe sieht etwa so aus:

{"data":[{"id":"google/gemma-4-26B-A4B-it","object":"model","owned_by":"organization-owner","permission":[],"metadata":{"model_id":"google/gemma-4-26B-A4B-it","input_modality":"text","max_request_context_length":8192}}],"object":"list"}

Senden Sie einen Prompt für die Chatvervollständigung an den Endpunkt /v1/chat/completions:

curl -X POST http://localhost:8000/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "google/gemma-4-26B-A4B-it",
        "messages": [
            {"role": "system", "content": "You are a helpful AI assistant."},
            {"role": "user", "content": "Explain in one sentence how tensor parallelism and expert parallelism accelerate Mixture-of-Experts LLM inference on TPUs."}
        ],
        "max_tokens": 128,
        "temperature": 0.0
    }'

Die Ausgabe auf Ironwood (TPU7x) sieht in etwa so aus:

{"id":"chatcmpl-ec30e61c-011d-49cd-989a-9ae9f59448a8","object":"chat.completion","created":1791887860,"model":"google/gemma-4-26B-A4B-it","choices":[{"index":0,"message":{"role":"assistant","content":"Tensor parallelism accelerates inference by partitioning large weight matrices across TPU cores to reduce per-device memory and computation, while expert parallelism distributes the specialized MoE layers across different devices to enable massive model scaling and parallelize the processing of different tokens through different experts.","refusal":null,"annotations":null,"audio":null,"function_call":null,"tool_calls":[],"reasoning":null},"logprobs":null,"finish_reason":"stop","stop_reason":106,"token_ids":null,"routed_experts":null}],"service_tier":null,"system_fingerprint":"vllm-0.22.1-tp16-ep-aaf76e1a","usage":{"prompt_tokens":49,"total_tokens":100,"completion_tokens":51,"prompt_tokens_details":null},"prompt_logprobs":null,"prompt_token_ids":null,"prompt_text":null,"kv_transfer_params":null}

Optional: Benchmarken Sie den Bereitstellungsdurchsatz und die Latenz mit dem Befehl vllm bench serve. Mit dem folgenden Befehl wird der HTTP-Lastgenerator-Client vllm bench serve von einem Ray-TPU-Worker-Pod ausgeführt und 64 gleichzeitige Anfragen an den Ray Serve-HTTP-Proxy auf Port 8000 gesendet:

WORKER_POD=$(kubectl get pods -l ray.io/node-type=worker -o jsonpath='{.items[0].metadata.name}')
kubectl exec -it ${WORKER_POD} -- vllm bench serve \
    --backend openai-chat \
    --endpoint /v1/chat/completions \
    --model google/gemma-4-26B-A4B-it \
    --dataset-name random \
    --random-input-len 1024 \
    --random-output-len 512 \
    --max-concurrency 32 \
    --num-prompts 64 \
    --num-warmups 4 \
    --temperature 0.0 \
    --ignore-eos \
    --host ray-serve-torchtpu-head-svc \
    --port 8000

Die Ausgabe sieht etwa so aus:

Trillium (v6e)

============ Serving Benchmark Result ============
Successful requests:                     64
Failed requests:                         0
Maximum request concurrency:             32
Benchmark duration (s):                  8.41
Total input tokens:                      66392
Total generated tokens:                  32768
Request throughput (req/s):              7.61
Output token throughput (tok/s):         3894.26
Peak output token throughput (tok/s):    4409.00
Peak concurrent requests:                64.00
Total token throughput (tok/s):          11784.50
---------------Time to First Token----------------
Mean TTFT (ms):                          414.57
Median TTFT (ms):                        348.94
P99 TTFT (ms):                           640.26
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          7.37
Median TPOT (ms):                        7.29
P99 TPOT (ms):                           8.24
---------------Inter-token Latency----------------
Mean ITL (ms):                           7.37
Median ITL (ms):                         0.00
P99 ITL (ms):                            25.94
==================================================

Ironwood (TPU7x)

============ Serving Benchmark Result ============
Successful requests:                     64
Failed requests:                         0
Maximum request concurrency:             32
Benchmark duration (s):                  9.65
Total input tokens:                      66392
Total generated tokens:                  32768
Request throughput (req/s):              6.64
Output token throughput (tok/s):         3397.33
Peak output token throughput (tok/s):    4574.00
Peak concurrent requests:                64.00
Total token throughput (tok/s):          10280.73
---------------Time to First Token----------------
Mean TTFT (ms):                          654.11
Median TTFT (ms):                        522.30
P99 TTFT (ms):                           1288.23
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          8.13
Median TPOT (ms):                        8.08
P99 TPOT (ms):                           9.69
---------------Inter-token Latency----------------
Mean ITL (ms):                           8.12
Median ITL (ms):                         9.31
P99 ITL (ms):                            26.85
==================================================

Deklarativ mit einer benutzerdefinierten RayService-Ressource bereitstellen

Für Produktionsbereitstellungen können Sie sowohl den Multi-Host-Ray-Cluster als auch die Ray Serve-LLM-Anwendung als eine einzige deklarative benutzerdefinierte RayService-Ressource verwalten, anstatt einen Job manuell einzureichen. Da der Ray Serve LLM-Builder die Modellkonfiguration direkt in der Serve-Konfiguration akzeptiert, benötigen Sie kein separates Python-Skript. KubeRay stellt den Cluster bereit, stellt die Serve-Anwendung bereit, überwacht den Zustand der Replikate und erstellt einen Kubernetes-ClusterIP-Dienst auf Port 8000, wenn die Anwendung bereit ist, Traffic zu verarbeiten.

Wenn die benutzerdefinierte RayCluster-Ressource aus dem vorherigen Abschnitt noch ausgeführt wird, beenden Sie die Portweiterleitung und löschen Sie die benutzerdefinierte RayCluster-Ressource, um das TPU-Slice freizugeben:
```
pkill -f "kubectl port-forward" || true
kubectl delete raycluster ray-serve-torchtpu --ignore-not-found
```

Sehen Sie sich das RayService-Manifest für Ihre TPU-Version im Beispiel-Repository an:

Trillium (v6e)

apiVersion: ray.io/v1
kind: RayService
metadata:
  name: ray-serve-torchtpu
spec:
  serveConfigV2: |
    applications:
    - name: default
      route_prefix: /
      import_path: ray.serve.llm:build_openai_app
      args:
        llm_configs:
        - model_loading_config:
            model_id: google/gemma-4-26B-A4B-it
          accelerator_type: TPU-V6E
          accelerator_config:
            kind: tpu
            topology: 4x4
          placement_group_config:
            bundle_per_worker:
              TPU: 1
          engine_kwargs:
            tensor_parallel_size: 16
            enable_expert_parallel: true
            enable_prefix_caching: true
            language_model_only: true
            disable_hybrid_kv_cache_manager: true
            max_model_len: 8192
            max_num_batched_tokens: 4096
          experimental_configs:
            stream_batching_interval_ms: 10
          runtime_env:
            env_vars:
              TPU_MULTIHOST_BACKEND: "ray"
              # Execute Gemma 4 routed expert layers on TPU TensorCores instead of SparseCore.
              USE_MOE_SPARSE_CORE: "0"
              VLLM_RAY_EXTRA_ENV_VARS_TO_COPY: "USE_MOE_SPARSE_CORE"
  rayClusterConfig:
    headGroupSpec:
      rayStartParams:
        num-cpus: "0"
      template:
        spec:
          containers:
          - name: ray-head
            image: ${IMAGE_URI}
            env:
            - name: HF_TOKEN
              valueFrom:
                secretKeyRef:
                  name: hf-secret
                  key: hf_api_token
            ports:
            - containerPort: 6379
              name: gcs
            - containerPort: 8265
              name: dashboard
            - containerPort: 10001
              name: client
            - containerPort: 8000
              name: serve
            resources:
              limits:
                cpu: "4"
                memory: "16Gi"
              requests:
                cpu: "4"
                memory: "16Gi"
    workerGroupSpecs:
    - groupName: v6e-16
      replicas: 1
      numOfHosts: 4
      template:
        spec:
          containers:
          - name: ray-worker
            image: ${IMAGE_URI}
            env:
            - name: HF_TOKEN
              valueFrom:
                secretKeyRef:
                  name: hf-secret
                  key: hf_api_token
            resources:
              limits:
                cpu: "64"
                memory: "200Gi"
                google.com/tpu: "4"
              requests:
                cpu: "64"
                memory: "200Gi"
                google.com/tpu: "4"
            volumeMounts:
            - mountPath: /dev/shm
              name: dshm
          volumes:
          - name: dshm
            emptyDir:
              medium: Memory
              sizeLimit: 100Gi
          nodeSelector:
            cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
            cloud.google.com/gke-tpu-topology: 4x4

Ironwood (TPU7x)

apiVersion: ray.io/v1
kind: RayService
metadata:
  name: ray-serve-torchtpu
spec:
  serveConfigV2: |
    applications:
    - name: default
      route_prefix: /
      import_path: ray.serve.llm:build_openai_app
      args:
        llm_configs:
        - model_loading_config:
            model_id: google/gemma-4-26B-A4B-it
          accelerator_type: TPU-V7X
          accelerator_config:
            kind: tpu
            topology: 2x2x2
          placement_group_config:
            bundle_per_worker:
              TPU: 1
          engine_kwargs:
            tensor_parallel_size: 16
            enable_expert_parallel: true
            enable_prefix_caching: true
            language_model_only: true
            disable_hybrid_kv_cache_manager: true
            max_model_len: 8192
            max_num_batched_tokens: 4096
          experimental_configs:
            stream_batching_interval_ms: 10
          runtime_env:
            env_vars:
              TPU_MULTIHOST_BACKEND: "ray"
              # Execute Gemma 4 routed expert layers on TPU TensorCores instead of SparseCore.
              USE_MOE_SPARSE_CORE: "0"
              VLLM_RAY_EXTRA_ENV_VARS_TO_COPY: "USE_MOE_SPARSE_CORE"
  rayClusterConfig:
    headGroupSpec:
      rayStartParams:
        num-cpus: "0"
      template:
        spec:
          containers:
          - name: ray-head
            image: ${IMAGE_URI}
            env:
            - name: HF_TOKEN
              valueFrom:
                secretKeyRef:
                  name: hf-secret
                  key: hf_api_token
            - name: RAY_TPU_RESOURCE_PER_CHIP
              value: "2"
            ports:
            - containerPort: 6379
              name: gcs
            - containerPort: 8265
              name: dashboard
            - containerPort: 10001
              name: client
            - containerPort: 8000
              name: serve
            resources:
              limits:
                cpu: "4"
                memory: "16Gi"
              requests:
                cpu: "4"
                memory: "16Gi"
    workerGroupSpecs:
    - groupName: tpu7x-16
      replicas: 1
      numOfHosts: 2
      resources:
        # Advertises 8 logical TPU cores (2 chiplets per physical chip) per host.
        TPU: "8"
      template:
        spec:
          containers:
          - name: ray-worker
            image: ${IMAGE_URI}
            env:
            - name: HF_TOKEN
              valueFrom:
                secretKeyRef:
                  name: hf-secret
                  key: hf_api_token
            - name: RAY_TPU_RESOURCE_PER_CHIP
              value: "2"
            resources:
              limits:
                cpu: "64"
                memory: "250Gi"
                google.com/tpu: "4"
              requests:
                cpu: "64"
                memory: "250Gi"
                google.com/tpu: "4"
            volumeMounts:
            - mountPath: /dev/shm
              name: dshm
          volumes:
          - name: dshm
            emptyDir:
              medium: Memory
              sizeLimit: 128Gi
          nodeSelector:
            cloud.google.com/gke-tpu-accelerator: tpu7x
            cloud.google.com/gke-tpu-topology: 2x2x2

Wenden Sie das RayService-Manifest an:

Trillium (v6e)

envsubst < ray-service.tpu-v6e-multihost.yaml | kubectl apply -f -

Ironwood (TPU7x)

envsubst < ray-service.tpu-v7x-multihost.yaml | kubectl apply -f -

Warten Sie, bis die benutzerdefinierte RayService-Ressource das Laden der Modellgewichtungen und das Kompilieren von XLA-Grafiken im TPU-Slice abgeschlossen hat und den Status Running erreicht:

kubectl get rayservice ray-serve-torchtpu
kubectl get svc ray-serve-torchtpu-serve-svc

Die Ausgabe sieht etwa so aus:

NAME                 SERVICE STATUS   NUM SERVE ENDPOINTS
ray-serve-torchtpu   Running          3
NAME                           TYPE        CLUSTER-IP       EXTERNAL-IP   PORT(S)    AGE
ray-serve-torchtpu-serve-svc   ClusterIP   34.118.239.213   <none>        8000/TCP   29s

Leiten Sie den lokalen Port 8000 an den ray-serve-torchtpu-serve-svc-Dienst weiter und fragen Sie das Modell mit den curl-Befehlen unter Bereitgestelltes Modell abfragen ab:
```
kubectl port-forward svc/ray-serve-torchtpu-serve-svc 8000:8000 >/dev/null 2>&1 &
```

Modellleistung beobachten

Informationen zum Überwachen der TPU-TensorCore-Auslastung, der HBM-Nutzung und der Ray-Cluster-Messwerte in Cloud Monitoring finden Sie unter Logs und Messwerte für Ray-Cluster in GKE erfassen und ansehen und TPU-Messwerte im Ray-Dashboard ansehen.

Auf das Ray-Dashboard zugreifen

So prüfen Sie den Status Ihrer Ray Serve-Bereitstellungen, Replikat-Akteure, Knotenlabels und Abfragelatenzen im integrierten Ray-Dashboard:

Wenn die Portweiterleitung zum Ray-Head-Dienst noch nicht ausgeführt wird, starten Sie sie im Hintergrund:
```
kubectl port-forward svc/ray-serve-torchtpu-head-svc 8265:8265 8000:8000 >/dev/null 2>&1 &
```
Klicken Sie in Cloud Shell auf Webvorschau und wählen Sie Vorschau auf Port 8265 aus oder öffnen Sie http://localhost:8265 in Ihrem lokalen Browser.
Klicken Sie auf die Tabs Bereitstellen, Cluster und Messwerte, um Ihre vLLM-Bereitstellungen, den Zustand der Replikate, automatisch erkannte ray.io/tpu-*-Knotenlabels, Messwerte zur Anfragelatenz und das Diagramm TPU-Auslastung nach Knoten aufzurufen.

Bereinigen

Wenn Sie vermeiden möchten, dass Ihrem Google Cloud Konto die in dieser Anleitung verwendeten Ressourcen in Rechnung gestellt werden, löschen Sie die benutzerdefinierte RayService- oder RayCluster-Ressource, den GKE-Cluster und das Artifact Registry-Repository:

Beenden Sie den Portweiterleitungsprozess im Hintergrund und löschen Sie die benutzerdefinierten Ray-Ressourcen und das Kubernetes-Secret:

pkill -f "kubectl port-forward" || true
kubectl delete rayservice ray-serve-torchtpu --ignore-not-found
kubectl delete raycluster ray-serve-torchtpu --ignore-not-found
kubectl delete secret hf-secret --ignore-not-found

Löschen Sie den GKE-Cluster:

Autopilot

gcloud container clusters delete ${CLUSTER_NAME} \
    --region=${REGION}

Standard

gcloud container clusters delete ${CLUSTER_NAME} \
    --zone=${ZONE}

Löschen Sie das Artifact Registry-Repository:

gcloud artifacts repositories delete ${ARTIFACT_REGISTRY} \
    --location=${REGION}

Nächste Schritte

Weitere Informationen zu PyTorch auf Cloud TPU
TorchTPU-Dokumentation lesen
Modelle auf TPUs mit mehreren Hosts mit Ray und TorchTPU optimieren
Offene Gemma-Modelle mithilfe von Multi-Host-TPUs in GKE mit Ray bereitstellen
Hier erfahren Sie, wie Sie Ray Serve-Bereitstellungen deklarativ mit einer benutzerdefinierten KubeRay RayService-Ressource verwalten.
TPUs in GKE
Ray Serve LLM-Dokumentation lesen

LLMs auf TPUs mit mehreren Hosts mit Ray und TorchTPU bereitstellen Mit Sammlungen den Überblick behalten Sie können Inhalte basierend auf Ihren Einstellungen speichern und kategorisieren.

Hintergrund

TPUs

vLLM und TorchTPU

Ray Serve-LLM in GKE

Ziele

Kosten

Hinweis

Umgebung vorbereiten

Benutzerdefiniertes Container-Image erstellen

GKE-Cluster erstellen

Autopilot

Standard

Trillium (v6e)

Ironwood (TPU7x)

Benutzerdefinierte RayCluster-Ressource mit mehreren Hosts bereitstellen

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Ray Serve-LLM-Anwendung erstellen und bereitstellen

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Bereitgestelltes Modell abfragen

Trillium (v6e)

Ironwood (TPU7x)

Deklarativ mit einer benutzerdefinierten RayService-Ressource bereitstellen

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Modellleistung beobachten

Auf das Ray-Dashboard zugreifen

Bereinigen

Autopilot

Standard

Nächste Schritte

LLMs auf TPUs mit mehreren Hosts mit Ray und TorchTPU bereitstellen