Serve an LLM using TPU Trillium on GKE with vLLM

Stay organized with collections Save and categorize content based on your preferences.

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:

gcloudcontainerclustersget-credentials${CLUSTER_NAME}--region=${REGION}

Create a Kubernetes Secret for Hugging Face credentials

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

   kubectlcreatenamespace${NAMESPACE}

2. Create a Kubernetes Secret that contains the Hugging Face token, run the following command:

   kubectlcreatesecretgenerichf-secret\--from-literal=hf_api_token=${HF_TOKEN}\--namespace${NAMESPACE}

Create a Cloud Storage bucket

In Cloud Shell, run the following command:

gcloudstoragebucketscreategs://${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:

   kubectlcreateserviceaccount${KSA_NAME}--namespace${NAMESPACE}

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

   gcloudstoragebucketsadd-iam-policy-bindinggs://${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:

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

Deploy vLLM model server

To deploy the vLLM model server, this tutorial uses a Kubernetes Deployment. A Deployment is a Kubernetes API object that lets you run multiple replicas of Pods that are distributed among the nodes in a cluster.

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

   apiVersion:apps/v1kind:Deploymentmetadata:name:vllm-tpuspec:replicas:1selector:matchLabels:app:vllm-tputemplate:metadata:labels:app:vllm-tpuannotations:gke-gcsfuse/volumes:"true"gke-gcsfuse/cpu-limit:"0"gke-gcsfuse/memory-limit:"0"gke-gcsfuse/ephemeral-storage-limit:"0"spec:serviceAccountName:KSA_NAMEnodeSelector:cloud.google.com/gke-tpu-topology:2x4cloud.google.com/gke-tpu-accelerator:tpu-v6e-slicecontainers:-name:vllm-tpuimage:docker.io/vllm/vllm-tpu:73aa7041bfee43581314e6f34e9a657137ecc092command:["python3","-m","vllm.entrypoints.openai.api_server"]args:---host=0.0.0.0---port=8000---tensor-parallel-size=8---max-model-len=4096---model=meta-llama/Llama-3.1-70B---download-dir=/data---max-num-batched-tokens=512---max-num-seqs=128env:-name:HUGGING_FACE_HUB_TOKENvalueFrom:secretKeyRef:name:hf-secretkey:hf_api_token-name:VLLM_XLA_CACHE_PATHvalue:"/data"-name:VLLM_USE_V1value:"1"ports:-containerPort:8000resources:limits:google.com/tpu:8readinessProbe:tcpSocket:port:8000initialDelaySeconds:15periodSeconds:10volumeMounts:-name:gcs-fuse-csi-ephemeralmountPath:/data-name:dshmmountPath:/dev/shmvolumes:-name:gke-gcsfuse-cacheemptyDir:medium:Memory-name:dshmemptyDir:medium:Memory-name:gcs-fuse-csi-ephemeralcsi:driver:gcsfuse.csi.storage.gke.iovolumeAttributes:bucketName:GSBUCKETmountOptions:"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:v1kind:Servicemetadata:name:vllm-servicespec:selector:app:vllm-tputype:LoadBalancerports:-name:httpprotocol:TCPport:8000targetPort:8000

2. Apply the manifest by running the following command:

   kubectlapply-fvllm-llama3-70b.yaml-n${NAMESPACE}

3. View the logs from the running model server:

   kubectllogs-f-lapp=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. To get the external IP address of the VLLM service, run the following command:

   exportvllm_service=$(kubectlgetservicevllm-service-ojsonpath='{.status.loadBalancer.ingress[0].ip}'-n${NAMESPACE})

2. Interact with the model using curl:

   curlhttp://$vllm_service:8000/v1/completions\ -H"Content-Type: application/json"\ -d'{ "model": "meta-llama/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/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:

   kubectlapply-fhttps://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:

   gcloudprojectsadd-iam-policy-bindingprojects/${PROJECT_ID}\--roleroles/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/v1kind:PodMonitoringmetadata:name:vllm-pod-monitoringspec:selector:matchLabels:app:vllm-tpuendpoints:-path:/metricsport:8000interval:15s

4. Apply it to the cluster:

   kubectlapply-fvllm_pod_monitor.yaml-n${NAMESPACE}

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/bashN=PARALLEL_PROCESSESexportvllm_service=$(kubectlgetservicevllm-service-ojsonpath='{.status.loadBalancer.ingress[0].ip}'-n${NAMESPACE})foriin$(seq1$N);dowhiletrue;docurlhttp://$vllm_service:8000/v1/completions-H"Content-Type: application/json"-d'{"model": "meta-llama/Llama-3.1-70B", "prompt": "Write a story about san francisco", "max_tokens": 1000, "temperature": 0}'done & # Run in the backgrounddonewait

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

2. Run the bash script:

   chmod+xload.sh 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:num_requests_waiting{cluster='CLUSTER_NAME'}

A line graph shows your vLLM metric (num_requests_waiting) measured over time. 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 example graph shows a starting pre-load value of 0, which reaches a maximum post-load value of close to 400 within one minute.

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.

In the following instructions, you deploy your HPA configuration by using the num_requests_waiting metric. For demonstration purposes, you set the metric to a low value so that the HPA configuration scales your vLLM replicas to two. 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/v2kind:HorizontalPodAutoscalermetadata:name:vllm-hpaspec:scaleTargetRef:apiVersion:apps/v1kind:Deploymentname:vllm-tpuminReplicas:1maxReplicas:2metrics:-type:Podspods:metric:name:prometheus.googleapis.com|vllm:num_requests_waiting|gaugetarget:type:AverageValueaverageValue: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:

   kubectlapply-fvllm-hpa.yaml-n${NAMESPACE}

   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:

   kubectlgethpa--watch-n${NAMESPACE}

   The output is similar to the following:

   NAME REFERENCE TARGETS MINPODS MAXPODS REPLICAS AGE vllm-hpa Deployment/vllm-tpu <unknown>/10 1 2 0 6s vllm-hpa Deployment/vllm-tpu 34972m/10 1 2 1 16s vllm-hpa Deployment/vllm-tpu 25112m/10 1 2 2 31s vllm-hpa Deployment/vllm-tpu 35301m/10 1 2 2 46s vllm-hpa Deployment/vllm-tpu 25098m/10 1 2 2 62s vllm-hpa Deployment/vllm-tpu 35348m/10 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 now.