This tutorial shows you how to deploy and serve
Llama 3.1 405B (nvidia/Llama-3.1-405B-Instruct-FP8),
a 405 billion-parameter instruction-tuned model quantized to FP8, on
Google Kubernetes Engine (GKE) on Google Cloud. The deployment uses
Cluster Toolkit, two
A4 (a4-highgpu-8g)
virtual machine (VM) instances with a total of 16 NVIDIA B200 GPUs,
NVIDIA Dynamo, and the
SGLang serving framework with
single-host prefill-decode (PD) disaggregation.
In the deployment configuration that's described in this tutorial, you use
Cluster Toolkit to provision a GKE cluster in
Standard mode. The cluster contains two a4-highgpu-8g nodes connected
over GPUDirect Remote Direct Memory Access (RDMA) networking. You then install
the KAI Scheduler and NVIDIA Dynamo platform with Grove gang scheduling.
Finally, you deploy a DynamoGraphDeployment custom resource that separates
prefill and decode across two single-host eight-GPU workers with tensor
parallelism and an FP8 KV cache.
Objectives
- Install Cluster Toolkit and provision a two-node A4 GKE cluster in Standard mode.
- Install the KAI Scheduler and NVIDIA Dynamo platform with Grove gang scheduling on your cluster.
- Deploy a PD-disaggregated SGLang inference server with one single-host A4 prefill worker and one single-host A4 decode worker.
- Send a chat completion request to the Dynamo frontend router.
- Run a serving benchmark workload.
- Clean up the deployed resources.
Costs
In this document, you use the following billable components of Google Cloud:
To generate a cost estimate based on your projected usage, use the pricing calculator.
Before you begin
Set up your environment for the tutorial:
- Install and initialize the
Google Cloud CLI, and install
kubectland Helm (version 3.10 or later). If you use Cloud Shell, these tools are preinstalled. - If you're running this tutorial from a Linux or macOS workstation rather than Cloud Shell, install the required Cluster Toolkit dependencies, such as Terraform and Git.
Enable the Kubernetes Engine API, Compute Engine API, Cloud Storage API, Identity and Access Management API, Cloud Resource Manager API, and Service Usage API in your Google Cloud project:
Get a Hugging Face access token
You need a Hugging Face access token to use this tutorial. You can sign up for a free account at Hugging Face. After you have an account, generate an access token:
- On the Welcome to Hugging Face page, click your account avatar and select Access tokens.
- On the Access tokens page, click Create new token.
- Select the Read token type and enter a name for your token.
- Your access token is displayed. Save the token in a safe place.
Required roles
To get the permissions that you need to complete this tutorial, ask your administrator to grant you the following IAM roles on your project:
- Kubernetes Engine Admin (
roles/container.admin) - Compute Admin (
roles/compute.admin) - Storage Admin (
roles/storage.admin) - Project IAM Admin (
roles/resourcemanager.projectIamAdmin) - Service Account Admin (
roles/iam.serviceAccountAdmin) - Service Account User (
roles/iam.serviceAccountUser) - Service Usage Consumer (
roles/serviceusage.serviceUsageConsumer)
For more information about granting roles, see Manage access to projects, folders, and organizations.
You might also be able to get the required permissions through custom roles or other predefined roles.
Set up your environment variables
Export the environment variables that you use throughout this tutorial:
Replace the following:
YOUR_PROJECT_ID: the ID of your Google Cloud project.YOUR_REGION: the region where your A4 capacity is located and where you want to deploy your cluster, such asus-central1.YOUR_ZONE: the zone where your A4 capacity is located, such asus-central1-b.YOUR_CLUSTER_NAME: the name for your GKE cluster and Cluster Toolkit deployment. Use a name of 20 characters or fewer so that the generated node pool service account ID (YOUR_CLUSTER_NAME-gke-np-sa) stays within the 30-character Identity and Access Management (IAM) limit.YOUR_RESERVATION_NAME: the name of your A4 reservation.YOUR_BUCKET_NAME: a globally unique name for the Cloud Storage bucket that stores the Terraform state for Cluster Toolkit.YOUR_HF_TOKEN: your Hugging Face access token.
Create the GKE cluster
Download Cluster Toolkit and provision a GKE cluster in
Standard mode with two a4-highgpu-8g nodes, GPUDirect RDMA
networking, and DRA support:
Sign in to the Google Cloud CLI and generate Application Default Credentials for Terraform:
Download and extract the prebuilt Cluster Toolkit bundle for the version specified in
CLUSTER_TOOLKIT_VERSION(v1.104.0):If you run this command on macOS or on an
arm64workstation, setOS="mac"orARCH="arm64"to match your platform.To store the Terraform state for your Cluster Toolkit deployment, create a Cloud Storage bucket with object versioning enabled:
Create the
cluster-toolkit/examples/gke-a4/gke-a4-deployment.yamldeployment configuration file for twoa4-highgpu-8gnodes:By default, the deployment file is configured to consume a specific reservation. If you don't have a specific reservation, then before you deploy the cluster, you can update
cluster-toolkit/examples/gke-a4/gke-a4-deployment.yamlto use an alternative consumption option, such as Spot VMs or Flex-start VMs. For configuration instructions, see GKE consumption options.Change to the
cluster-toolkitdirectory, deploy the GKE cluster by using thegcluster deploycommand, and then return to your working directory:Cluster creation takes approximately 20 to 30 minutes. If you run
./gcluster deploywithout the--auto-approveflag, then to apply the proposed changes, enterawhen you're prompted to.Configure
kubectlcredentials to communicate with your cluster:
Install KAI Scheduler and NVIDIA Dynamo
Prepare your cluster by creating a Hugging Face Secret and installing the KAI Scheduler and NVIDIA Dynamo platform with Grove gang scheduling:
Create a Kubernetes Secret named
hf-secretto store your Hugging Face access token:Download the NVIDIA Dynamo platform Helm chart (
v1.5.0), install the KAI Scheduler, and install the Dynamo platform with Grove enabled:
Deploy the single-host disaggregated SGLang server
The workload manifest (llama3.1-405b-singlehost-a4.yaml) defines a Kubernetes
ResourceClaimTemplate (all-mrdma), which allocates the node's multi-NIC RDMA
devices (mrdma.google.com) through DRA, and an NVIDIA Dynamo
DynamoGraphDeployment (llama3-405b-a4-disagg) with Grove gang scheduling
enabled. The deployment consists of the following:
- Frontend (
Frontend): runs one router replica that accepts client requests, discovers registered prefill and decode workers, and coordinates disaggregated request routing. - Prefill worker (
Prefill): runs a single-host worker on onea4-highgpu-8gnode with eight NVIDIA B200 GPUs. On startup, the Pod sources the NCCL/gIB environment, downloads thenvidia/Llama-3.1-405B-Instruct-FP8checkpoint from Hugging Face Hub to the node's Local NVMe SSD, and launchespython3 -m dynamo.sglangin--disaggregation-mode prefillwith the following flags:--tensor-parallel-size 8--kv-cache-dtype fp8_e4m3--context-length 16384--mem-fraction-static 0.85
- Decode worker (
Decode): runs a single-host worker on the seconda4-highgpu-8gnode with eight NVIDIA B200 GPUs. On startup, the Pod caches the model weights on the node's Local NVMe SSD and launchespython3 -m dynamo.sglangin--disaggregation-mode decodewith the following flags:--tensor-parallel-size 8--kv-cache-dtype fp8_e4m3--context-length 16384--mem-fraction-static 0.85
Deploy and monitor the disaggregated inference server:
Download the
llama3.1-405b-singlehost-a4.yamlworkload manifest:Alternatively, create a file named
llama3.1-405b-singlehost-a4.yamlwith the following content:Apply the workload manifest to your cluster:
To monitor weight download and server initialization, list the Pods in your cluster, and stream the logs of the prefill and decode Pods:
Replace
YOUR_POD_NAMEwith the name of the prefill Pod or decode Pod from thekubectl get podsoutput.On the first run, it takes approximately 10 to 20 minutes to download the 405B FP8 model weights and initialize the SGLang servers. Before you send any requests, wait until the frontend, prefill, and decode Pods show
1/1in theREADYcolumn ofkubectl get pods:NAME READY STATUS RESTARTS AGE llama3-405b-a4-disagg-0-decode-tgl7q 1/1 Running 0 15m llama3-405b-a4-disagg-0-frontend-wrxq5 1/1 Running 0 15m llama3-405b-a4-disagg-0-prefill-67h5s 1/1 Running 0 15m
Send inference requests
Send a chat completion request to the deployed Llama 3.1 405B service:
Forward local port
8000to thellama3-405b-a4-disagg-frontendservice:Leave port forwarding running in this terminal session, or run the command in the background.
In a second terminal session, send a chat completion request to the
/v1/chat/completionsendpoint:
Run a benchmark workload
Benchmark the deployed Llama 3.1 405B server from the Dynamo frontend Pod by
using sglang.bench_serving. Configure the benchmark for 1,024 input tokens,
1,024 output tokens (1k/1k), 2,048 prompts, and a maximum concurrency of 512:
After the benchmark completes, look for output that contains metrics such as the following:
Output token throughputTotal token throughputTTFTTPOT
Clean up
To avoid incurring additional ongoing charges to your Google Cloud account for the resources used in this tutorial, delete the Kubernetes workloads, the GKE cluster, and your Cloud Storage bucket:
Delete the DynamoGraphDeployment and ResourceClaimTemplate resources:
Uninstall the NVIDIA Dynamo platform and KAI Scheduler Helm releases:
Delete the Hugging Face Kubernetes Secret:
From the
cluster-toolkitdirectory, destroy the GKE cluster and its associated infrastructure by using thegcluster destroycommand:To delete the Cloud Storage bucket that stores the Terraform state, use the
gcloud storage rmcommand:gcloud storage rm --recursive "gs://${GCS_BUCKET}"
What's next
- Learn more about NVIDIA Dynamo.
- Learn more about SGLang.
- Learn more about A4 VMs with NVIDIA B200 GPUs.
- Learn more about Cluster Toolkit.
- Learn more about managing AI-optimized GKE clusters.