Serve an LLM using TPU Trillium (v6e) on GKE with vLLM

Create a GKE cluster

You can serve LLMs on TPUs in a GKE Autopilot or Standard cluster. We recommend that you use a Autopilot cluster for a fully managed Kubernetes experience. To choose the GKE mode of operation that's the best fit for your workloads, see Choose a GKE mode of operation.

Configure kubectl to communicate with your cluster

To configure kubectl to communicate with your cluster, run the following command:

gcloud container clusters get-credentials CLUSTER_NAME --location=ZONE

Create a Kubernetes Secret for Hugging Face credentials

To create a Kubernetes Secret that contains the Hugging Face token, run the following command:

kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=HUGGING_FACE_TOKEN \
    --namespace NAMESPACE

Create a Cloud Storage bucket

In Cloud Shell, run the following command:

gcloud storage buckets create gs://GSBUCKET \
    --uniform-bucket-level-access

This creates a Cloud Storage bucket to store the model files you download from Hugging Face.

Set up a Kubernetes ServiceAccount to access the bucket

1. Create a namespace to use for the Kubernetes ServiceAccount. You can skip this step if you are using the default namespace:

   kubectl create namespace NAMESPACE
   

2. Create the Kubernetes ServiceAccount:

   kubectl create serviceaccount KSA_NAME --namespace NAMESPACE
   

3. Grant read-write access to the Kubernetes ServiceAccount in order to access the Cloud Storage bucket:

   gcloud storage buckets add-iam-policy-binding gs://GSBUCKET \
     --member "principal://iam.googleapis.com/projects/PROJECT_NUMBER/locations/global/workloadIdentityPools/PROJECT_ID.svc.id.goog/subject/ns/NAMESPACE/sa/KSA_NAME" \
     --role "roles/storage.objectUser"
   

4. Alternatively, you can grant read-write access to all Cloud Storage buckets in the project:

   gcloud projects add-iam-policy-binding PROJECT_ID \
   --member "principal://iam.googleapis.com/projects/PROJECT_NUMBER/locations/global/workloadIdentityPools/PROJECT_ID.svc.id.goog/subject/ns/NAMESPACE/sa/KSA_NAME" \
   --role "roles/storage.objectUser"
   

   GKE creates the following resources for the LLM:

   1. A Cloud Storage bucket to store the downloaded model and the compilation cache. A Cloud Storage FUSE CSI driver reads the content of the bucket.
   2. Volumes with file caching enabled and the parallel download feature of Cloud Storage FUSE.
   Best practice:

   Use a file cache backed by tmpfs or Hyperdisk / Persistent Disk depending on the expected size of the model contents, for example, weight files. In this tutorial, you use Cloud Storage FUSE file cache backed by RAM.

Build and deploy the TPU image

Containerize your vLLM server:

1. Clone the vLLM repository and build the image:

   git clone https://github.com/vllm-project/vllm && cd vllm && git reset --hard cd34029e91ad2d38a58d190331a65f9096c0b157 && docker build -f Dockerfile.tpu . -t vllm-tpu
   

2. Push the image to Artifact Registry:

   gcloud artifacts repositories create vllm-tpu --repository-format=docker --location=REGION_NAME && \
   gcloud auth configure-docker REGION_NAME-docker.pkg.dev && \
   docker image tag vllm-tpu REGION_NAME-docker.pkg.dev/PROJECT_ID/vllm-tpu/vllm-tpu:latest && \
   docker push REGION_NAME-docker.pkg.dev/PROJECT_ID/vllm-tpu/vllm-tpu:latest
   

Deploy vLLM model server

To deploy the vLLM model server, follow these steps:

1. Inspect the manifest vllm-llama3-70b.yaml:

   
   apiVersion: apps/v1
   kind: Deployment
   metadata:
     name: vllm-tpu
   spec:
     replicas: 1
     selector:
       matchLabels:
         app: vllm-tpu
     template:
       metadata:
         labels:
           app: vllm-tpu
         annotations:
           gke-gcsfuse/volumes: "true"
           gke-gcsfuse/cpu-limit: "0"
           gke-gcsfuse/memory-limit: "0"
           gke-gcsfuse/ephemeral-storage-limit: "0"
       spec:
         serviceAccountName: KSA_NAME
         nodeSelector:
           cloud.google.com/gke-tpu-topology: 2x4
           cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
         containers:
         - name: vllm-tpu
           image: REGION_NAME-docker.pkg.dev/PROJECT_ID/vllm-tpu/vllm-tpu:latest
           command: ["python3", "-m", "vllm.entrypoints.openai.api_server"]
           args:
           - --host=0.0.0.0
           - --port=8000
           - --tensor-parallel-size=8
           - --max-model-len=8192
           - --model=meta-llama/Meta-Llama-3.1-70B
           - --download-dir=/data
           env: 
           - name: HUGGING_FACE_HUB_TOKEN
             valueFrom:
               secretKeyRef:
                 name: hf-secret
                 key: hf_api_token
           - name: VLLM_XLA_CACHE_PATH
             value: "/data"
           ports:
           - containerPort: 8000
           resources:
             limits:
               google.com/tpu: 8
           readinessProbe:
             tcpSocket:
               port: 8000
             initialDelaySeconds: 15
             periodSeconds: 10
           volumeMounts:
           - name: gcs-fuse-csi-ephemeral
             mountPath: /data
           - name: dshm
             mountPath: /dev/shm
         volumes:
         - name: gke-gcsfuse-cache
           emptyDir:
             medium: Memory
         - name: dshm
           emptyDir:
             medium: Memory
         - name: gcs-fuse-csi-ephemeral
           csi:
             driver: gcsfuse.csi.storage.gke.io
             volumeAttributes:
               bucketName: GSBUCKET
               mountOptions: "implicit-dirs,file-cache:enable-parallel-downloads:true,file-cache:parallel-downloads-per-file:100,file-cache:max-parallel-downloads:-1,file-cache:download-chunk-size-mb:10,file-cache:max-size-mb:-1"
   ---
   apiVersion: v1
   kind: Service
   metadata:
     name: vllm-service
   spec:
     selector:
       app: vllm-tpu
     type: LoadBalancer	
     ports:
       - name: http
         protocol: TCP
         port: 8000  
         targetPort: 8000
   

2. Apply the manifest by running the following command:

   kubectl apply -f vllm-llama3-70b.yaml
   

3. View the logs from the running model server:

   kubectl logs -f -l app=vllm-tpu
   

   The output should look similar to the following:

   INFO 09-20 19:03:48 launcher.py:19] Available routes are:
   INFO 09-20 19:03:48 launcher.py:27] Route: /openapi.json, Methods: GET, HEAD
   INFO 09-20 19:03:48 launcher.py:27] Route: /docs, Methods: GET, HEAD
   INFO 09-20 19:03:48 launcher.py:27] Route: /docs/oauth2-redirect, Methods: GET, HEAD
   INFO 09-20 19:03:48 launcher.py:27] Route: /redoc, Methods: GET, HEAD
   INFO 09-20 19:03:48 launcher.py:27] Route: /health, Methods: GET
   INFO 09-20 19:03:48 launcher.py:27] Route: /tokenize, Methods: POST
   INFO 09-20 19:03:48 launcher.py:27] Route: /detokenize, Methods: POST
   INFO 09-20 19:03:48 launcher.py:27] Route: /v1/models, Methods: GET
   INFO 09-20 19:03:48 launcher.py:27] Route: /version, Methods: GET
   INFO 09-20 19:03:48 launcher.py:27] Route: /v1/chat/completions, Methods: POST
   INFO 09-20 19:03:48 launcher.py:27] Route: /v1/completions, Methods: POST
   INFO 09-20 19:03:48 launcher.py:27] Route: /v1/embeddings, Methods: POST
   INFO:     Started server process [1]
   INFO:     Waiting for application startup.
   INFO:     Application startup complete.
   INFO:     Uvicorn running on http://0.0.0.0:8080 (Press CTRL+C to quit)
   (RayWorkerWrapper pid=25987) INFO 09-20 19:03:46 tpu_model_runner.py:290] Compilation for decode done in 202.93 s.
   (RayWorkerWrapper pid=25987) INFO 09-20 19:03:46 tpu_model_runner.py:283] batch_size: 256, seq_len: 1 [repeated 7x across cluster]
   INFO 09-20 19:03:53 metrics.py:351] Avg prompt throughput: 0.0 tokens/s, Avg generation throughput: 0.0 tokens/s, Running: 0 reqs, Swapped: 0 reqs, Pending: 0 reqs, GPU KV cache usage: 0.0%, CPU KV cache usage: 0.0%.
   

Serve the model

1. Run the following command to get the external IP address of the Service:

   export vllm_service=$(kubectl get service vllm-service -o jsonpath='{.status.loadBalancer.ingress[0].ip}')
   

2. In a new terminal, interact with the model using curl:

   curl http://$vllm_service:8000/v1/completions \
   -H "Content-Type: application/json" \
   -d '{
       "model": "meta-llama/Meta-Llama-3.1-70B",
       "prompt": "San Francisco is a",
       "max_tokens": 7,
       "temperature": 0
   }'
   

   The output should be similar to the following:

   {"id":"cmpl-6b4bb29482494ab88408d537da1e608f","object":"text_completion","created":1727822657,"model":"meta-llama/Meta-Llama-3-8B","choices":[{"index":0,"text":" top holiday destination featuring scenic beauty and","logprobs":null,"finish_reason":"length","stop_reason":null,"prompt_logprobs":null}],"usage":{"prompt_tokens":5,"total_tokens":12,"completion_tokens":7}}
   

Set up the custom autoscaler

In this section, you set up horizontal Pod autoscaling using custom Prometheus metrics. You use the Google Cloud Managed Service for Prometheus metrics from the vLLM server.

To learn more, see Google Cloud Managed Service for Prometheus. This should be enabled by default on the GKE cluster.

1. Set up the Custom Metrics Stackdriver Adapter on your cluster:

   kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/k8s-stackdriver/master/custom-metrics-stackdriver-adapter/deploy/production/adapter_new_resource_model.yaml
   

2. Add the Monitoring Viewer role to the service account that the Custom Metrics Stackdriver Adapter uses:

   gcloud projects add-iam-policy-binding projects/PROJECT_ID \
       --role roles/monitoring.viewer \
       --member=principal://iam.googleapis.com/projects/PROJECT_NUMBER/locations/global/workloadIdentityPools/PROJECT_ID.svc.id.goog/subject/ns/custom-metrics/sa/custom-metrics-stackdriver-adapter
   

3. Save the following manifest as vllm_pod_monitor.yaml:

   
   apiVersion: monitoring.googleapis.com/v1
   kind: PodMonitoring
   metadata:
    name: vllm-pod-monitoring
   spec:
    selector:
      matchLabels:
        app: vllm-tpu
    endpoints:
    - path: /metrics
      port: 8000
      interval: 15s
   

4. Apply it to the cluster:

   kubectl apply -f vllm_pod_monitor.yaml
   

Create load on the vLLM endpoint

Create load to the vLLM server to test how GKE autoscales with a custom vLLM metric.

1. Run a bash script (load.sh) to send N number of parallel requests to the vLLM endpoint:

   #!/bin/bash
   N=PARALLEL_PROCESSES
   export vllm_service=$(kubectl get service vllm-service -o jsonpath='{.status.loadBalancer.ingress[0].ip}')
   for i in $(seq 1 $N); do
     while true; do
       curl http://$vllm_service:8000/v1/completions -H "Content-Type: application/json" -d '{"model": "meta-llama/Meta-Llama-3.1-70B", "prompt": "Write a story about san francisco", "max_tokens": 100, "temperature": 0}'
     done &  # Run in the background
   done
   wait
   

   Replace PARALLEL_PROCESSES with the number of parallel processes that you want to run.

2. Run the bash script:

   nohup ./load.sh &
   

Verify that Google Cloud Managed Service for Prometheus ingests the metrics

After Google Cloud Managed Service for Prometheus scrapes the metrics and you're adding load to the vLLM endpoint, you can view metrics on Cloud Monitoring.

1. In the Google Cloud console, go to the Metrics explorer page.

   Go to Metrics explorer

2. Click < > PromQL.

3. Enter the following query to observe traffic metrics:

   vllm:avg_generation_throughput_toks_per_s{cluster='CLUSTER_NAME_HERE'}
   

In the line graph, the vLLM metric scales up from 0 (pre-load) to a value (post-load). This graph confirms your vLLM metrics are being ingested into Google Cloud Managed Service for Prometheus.

The following image is an example of a graph after the load script execution. In this case, the model server is serving around 2,000 generation tokens per second.

Deploy the Horizontal Pod Autoscaler configuration

In this section, you deploy the Horizontal Pod Autoscaler configuration.

1. Save the following manifest as vllm-hpa.yaml:

   
   apiVersion: autoscaling/v2
   kind: HorizontalPodAutoscaler
   metadata:
    name: vllm-hpa
   spec:
    scaleTargetRef:
      apiVersion: apps/v1
      kind: Deployment
      name: vllm-tpu
    minReplicas: 1
    maxReplicas: 2
    metrics:
      - type: Pods
        pods:
          metric:
            name: prometheus.googleapis.com|vllm:num_requests_waiting|gauge
          target:
            type: AverageValue
            averageValue: 1
   

   The vLLM metrics in Google Cloud Managed Service for Prometheus follow the vllm:metric_name format.

   Best practice:

   Use num_requests_waiting for scaling throughput. Use gpu_cache_usage_perc for latency-sensitive TPU use cases.

2. Deploy the Horizontal Pod Autoscaler configuration:

   kubectl apply -f vllm-hpa.yaml
   

   GKE schedules another Pod to deploy, which triggers the node pool autoscaler to add a second node before it deploys the second vLLM replica.

3. Watch the progress of the Pod autoscaling:

   kubectl get hpa --watch
   

   The output is similar to the following:

   NAME       REFERENCE             TARGETS       MINPODS   MAXPODS   REPLICAS   AGE
   vllm-hpa   Deployment/vllm-tpu   <unknown>/5   1         2         0          6s
   vllm-hpa   Deployment/vllm-tpu   34972m/5      1         2         1          16s
   vllm-hpa   Deployment/vllm-tpu   25112m/5      1         2         2          31s
   vllm-hpa   Deployment/vllm-tpu   35301m/5      1         2         2          46s
   vllm-hpa   Deployment/vllm-tpu   25098m/5      1         2         2          62s
   vllm-hpa   Deployment/vllm-tpu   35348m/5      1         2         2          77s
   

4. Wait for 10 minutes and repeat the steps in the Verify that Google Cloud Managed Service for Prometheus ingests the metrics section. Google Cloud Managed Service for Prometheus ingests the metrics from both vLLM endpoints.