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

1. Create a namespace. You can skip this step if you are using the default namespace:

   kubectl create namespace NAMESPACE
   

2. 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 the Kubernetes ServiceAccount:

   kubectl create serviceaccount KSA_NAME --namespace NAMESPACE
   

2. 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"
   

3. 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.

(Optional) Build and deploy the TPU image

If you need finer-grain control over the contents of your Docker image, choose this option.

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 2e33fe419186c65a18da6668972d61d7bbc31564 && 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 Deployment manifest saved as vllm-llama3-70b.yaml. A Deployment is a Kubernetes API object that lets you run multiple replicas of Pods that are distributed among the nodes in a cluster:

   
   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: IMAGE_NAME
           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 -n NAMESPACE
   

3. View the logs from the running model server:

   kubectl logs -f -l app=vllm-tpu -n NAMESPACE
   

   The output should look similar to the following:

   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)
   

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

When deciding which metric to autoscale on, we recommend the following metrics for vLLM TPU:

• num_requests_waiting: This metric relates to the number of requests waiting in the model server's queue. This number starts to noticeably grow when the kv cache is full.

• gpu_cache_usage_perc: This metric relates to the kv cache utilization, which directly correlates to the number of requests being processed for a given inference cycle on the model server. Note that this metric works the same on GPUs and TPUs, though it is tied to the GPU naming schema.

We recommend that you use num_requests_waiting when optimizing for throughput and cost, and when your latency targets are achievable with your model server's maximum throughput.

We recommend that you use gpu_cache_usage_perc when you have latency-sensitive workloads where queue-based scaling isn't fast enough to meet your requirements.

For further explanation, check out Best practices for autoscaling large language model (LLM) inference workloads with TPUs.

When selecting an averageValue target for your HPA config, you will have to determine this experimentally. Check out the Save on GPUs: Smarter autoscaling for your GKE inferencing workloads blog post for additional ideas on how to optimize this part. The profile-generator used in this blog post works for vLLM TPU as well.

To deploy the Horizontal Pod Autoscaler configuration using num_requests_waiting, follow these steps:

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: 10
   

   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.