Serve an LLM using TPU Trillium 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 = 
 ${ 
 CONTROL_PLANE_LOCATION 
 } 
 

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 
    = 
    ${ 
    HF_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.

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 following Deployment manifest saved as vllm-llama3-70b.yaml , that uses a single replica:

     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 
    : 
     
    docker.io/vllm/vllm-tpu:73aa7041bfee43581314e6f34e9a657137ecc092 
     
    command 
    : 
     
    [ 
    "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=128 
     
    env 
    : 
     
     
    - 
     
    name 
    : 
     
    HUGGING_FACE_HUB_TOKEN 
     
    valueFrom 
    : 
     
    secretKeyRef 
    : 
     
    name 
    : 
     
    hf-secret 
     
    key 
    : 
     
    hf_api_token 
     
    - 
     
    name 
    : 
     
    VLLM_XLA_CACHE_PATH 
     
    value 
    : 
     
    "/data" 
     
    - 
     
    name 
    : 
     
    VLLM_USE_V1 
     
    value 
    : 
     
    "1" 
     
    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 
    
   

   If you scale up the Deployment to multiple replicas, the concurrent writes to the VLLM_XLA_CACHE_PATH will cause the error: RuntimeError: filesystem error: cannot create directories . To prevent this error, you have two options:

   1. Remove the XLA cache location by removing the following block from the Deployment YAML. This means all replicas will recompile the cache.

       - name: VLLM_XLA_CACHE_PATH
        value: "/data" 
      

   2. Scale the Deployment to 1 , and wait for the first replica to become ready and write to the XLA cache. Then scale to additional replicas. This allows the remainder of the replicas to read the cache, without attempting to write it.

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

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

2. Interact with the model using curl :

    curl  
   http:// $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:

    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  
   -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/bash 
    N 
    = 
    PARALLEL_PROCESSES 
    export 
     
    vllm_service 
    = 
    $( 
   kubectl  
   get  
   service  
   vllm-service  
   -o  
    jsonpath 
    = 
    '{.status.loadBalancer.ingress[0].ip}' 
     
   -n  
    ${ 
    NAMESPACE 
    } 
    ) 
    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/Llama-3.1-70B", "prompt": "Write a story about san francisco", "max_tokens": 1000, "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:

    chmod  
   +x  
   load.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 explorerpage.

   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_waitingmetric. 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/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  
   -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:

    kubectl  
   get  
   hpa  
   --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.