在 A4 VM 上使用 Dynamo 和 SGLang，透過 PD 解構在 GKE 提供 Llama 3.1 405B

設定環境變數

匯出本教學課程中使用的環境變數：

export PROJECT_ID="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export RESERVATION_NAME="YOUR_RESERVATION_NAME"
export GCS_BUCKET="YOUR_BUCKET_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export CLUSTER_TOOLKIT_VERSION="v1.104.0"
export DYNAMO_RELEASE_VERSION="1.5.0"

gcloud config set project "${PROJECT_ID}"
gcloud config set billing/quota_project "${PROJECT_ID}"

更改下列內容：

• YOUR_PROJECT_ID： Google Cloud專案的 ID。
• YOUR_REGION：A4 容量所在及要部署叢集的區域，例如 us-central1。
• YOUR_ZONE：A4 運算資源所在的可用區，例如 us-central1-b。
• YOUR_CLUSTER_NAME：GKE 叢集和 Cluster Toolkit 部署作業的名稱。請使用 20 個字元以內的名稱，確保產生的節點集區服務帳戶 ID (YOUR_CLUSTER_NAME-gke-np-sa) 不會超過 Identity and Access Management (IAM) 的 30 個字元限制。
• YOUR_RESERVATION_NAME：A4 預留項目的名稱。
• YOUR_BUCKET_NAME：Cloud Storage bucket 的全域不重複名稱，用於儲存 Cluster Toolkit 的 Terraform 狀態。
• YOUR_HF_TOKEN：您的 Hugging Face 存取權杖。

建立 GKE 叢集

下載 Cluster Toolkit，並以標準模式佈建 GKE 叢集，其中包含兩個 a4-highgpu-8g 節點、GPUDirect RDMA 網路和 DRA 支援：

1. 登入 Google Cloud CLI，並為 Terraform 產生應用程式預設憑證：

   gcloud auth login
   gcloud auth application-default login

2. 下載並解壓縮預先建構的 Cluster Toolkit 套件，版本請依 CLUSTER_TOOLKIT_VERSION (v1.104.0) 指定：

   TAG="${CLUSTER_TOOLKIT_VERSION:-v1.104.0}"
   OS="linux"
   ARCH="amd64"
   mkdir -p cluster-toolkit && curl -L "https://github.com/GoogleCloudPlatform/cluster-toolkit/releases/download/${TAG}/gcluster_bundle_${OS}_${ARCH}.tgz" | tar -xz -C cluster-toolkit

   如果您在 macOS 或 arm64 工作站上執行這個指令，請將 OS="mac" 或 ARCH="arm64" 設為與平台相符。

3. 如要儲存 Cluster Toolkit 部署作業的 Terraform 狀態，請建立 Cloud Storage bucket，並啟用物件版本管理功能：

   gcloud storage buckets create "gs://${GCS_BUCKET}" \
       --default-storage-class=STANDARD \
       --project="${PROJECT_ID}" \
       --location="${REGION}" \
       --uniform-bucket-level-access
   gcloud storage buckets update "gs://${GCS_BUCKET}" --versioning

4. 為兩個 a4-highgpu-8g 節點建立 cluster-toolkit/examples/gke-a4/gke-a4-deployment.yaml 部署設定檔：

   cat <<EOF > cluster-toolkit/examples/gke-a4/gke-a4-deployment.yaml
   terraform_backend_defaults:
     type: gcs
     configuration:
       bucket: ${GCS_BUCKET}
   
   vars:
     project_id: ${PROJECT_ID}
     deployment_name: ${CLUSTER_NAME}
     region: ${REGION}
     zone: ${ZONE}
     static_node_count: 2
     authorized_cidr: 0.0.0.0/0
     reservation_affinity:
       consume_reservation_type: SPECIFIC_RESERVATION
       specific_reservations:
       - name: ${RESERVATION_NAME}
   EOF

   根據預設，部署檔案會設為使用特定預留項目。如果沒有特定預留項目，請先更新 cluster-toolkit/examples/gke-a4/gke-a4-deployment.yaml，改用其他消耗選項 (例如 Spot VM 或彈性啟動型 VM)，再部署叢集。如需設定操作說明，請參閱「GKE 消耗選項」。

5. 切換至 cluster-toolkit 目錄，使用 gcluster deploy 指令部署 GKE 叢集，然後返回目前使用的目錄：

   cd cluster-toolkit
   ./gcluster deploy \
       -d examples/gke-a4/gke-a4-deployment.yaml \
       examples/gke-a4/gke-a4.yaml \
       --auto-approve -w
   cd ..

   建立叢集約需 20 到 30 分鐘。如果您執行 ./gcluster deploy 時未使用 --auto-approve 旗標，請在系統提示時輸入 a，套用建議的變更。

6. 設定 kubectl 憑證，以便與叢集通訊：

   gcloud container clusters get-credentials "${CLUSTER_NAME}" \
       --location="${REGION}" \
       --project="${PROJECT_ID}"

安裝 KAI Scheduler 和 NVIDIA Dynamo

建立 Hugging Face Secret，並安裝 KAI 排程器和 NVIDIA Dynamo 平台 (採用 Grove Gang 排程)，為叢集做好準備：

1. 建立名為 hf-secret 的 Kubernetes Secret，用於儲存 Hugging Face 存取權杖：

   kubectl create secret generic hf-secret \
       --from-literal=hf_api_token="${HF_TOKEN}" \
       --dry-run=client -o yaml | kubectl apply -f -

2. 下載 NVIDIA Dynamo 平台 Helm 圖表 (v1.5.0)，安裝 KAI 排程器，並安裝已啟用 Grove 的 Dynamo 平台：

   RELEASE_VERSION="${DYNAMO_RELEASE_VERSION:-1.5.0}"
   helm fetch "https://helm.ngc.nvidia.com/nvidia/ai-dynamo/charts/dynamo-platform-${RELEASE_VERSION}.tgz"
   tar -zxf "dynamo-platform-${RELEASE_VERSION}.tgz"
   
   helm upgrade --install kai-scheduler ./dynamo-platform/charts/kai-scheduler \
       --set scheduler.args.default-staleness-grace-period="-1s" \
       --set topologyMigration.enabled=false
   
   kubectl wait --for=condition=Available deployment/kai-operator --timeout=120s
   until kubectl get deployment/queue-controller >/dev/null 2>&1; do sleep 2; done
   kubectl wait --for=condition=Available deployment/queue-controller --timeout=120s
   
   helm upgrade --install dynamo-platform ./dynamo-platform \
       --set global.grove.install=true \
       --set global.grove.enabled=true \
       --set global.kai-scheduler.enabled=true

部署單一主機解構式 SGLang 伺服器

工作負載資訊清單 (llama3.1-405b-singlehost-a4.yaml) 定義了 Kubernetes ResourceClaimTemplate (all-mrdma)，透過 DRA 分配節點的多 NIC RDMA 裝置 (mrdma.google.com)，以及啟用 Grove 群組排程的 NVIDIA Dynamo DynamoGraphDeployment (llama3-405b-a4-disagg)。部署作業包含下列項目：

• 前端 (Frontend)：執行一個路由器副本，接受用戶端要求、探索已註冊的預先填入和解碼 worker，並協調分散式要求轉送。
• 預先填入 worker (Prefill)：在一個 a4-highgpu-8g 節點上執行單一主機 worker，並配備八個 NVIDIA B200 GPU。啟動時，Pod 會取得 NCCL/gIB 環境、從 Hugging Face Hub 下載 nvidia/Llama-3.1-405B-Instruct-FP8 檢查點至節點的本機 NVMe SSD，並使用下列旗標在 --disaggregation-mode prefill 中啟動 python3 -m dynamo.sglang：
  • --tensor-parallel-size 8
  • --kv-cache-dtype fp8_e4m3
  • --context-length 16384
  • --mem-fraction-static 0.85
• 解碼 worker (Decode)：在第二個 a4-highgpu-8g 節點上執行單一主機 worker，並配備八個 NVIDIA B200 GPU。啟動時，Pod 會將模型權重快取至節點的本機 NVMe SSD，並在 --disaggregation-mode decode 中啟動 python3 -m dynamo.sglang，並使用下列旗標：
  • --tensor-parallel-size 8
  • --kv-cache-dtype fp8_e4m3
  • --context-length 16384
  • --mem-fraction-static 0.85

部署及監控分解推論伺服器：

1. 下載 llama3.1-405b-singlehost-a4.yaml 工作負載資訊清單：

   curl -O https://raw.githubusercontent.com/GoogleCloudSamples/AIHypercomputerSamples/refs/heads/main/gpu/inference/llama3-405b/llama3.1-405b-singlehost-a4.yaml

   或者，您也可以建立名為 llama3.1-405b-singlehost-a4.yaml 的檔案，並加入下列內容：

   # ==============================================================================
   # Single-Host Disaggregated Serving of Llama 3.1 405B FP8 on GKE
   # using NVIDIA Dynamo Platform (v1.5.0) and SGLang on GPU A4 (2 x a4-highgpu-8g)
   # Topology: 1-node Prefill (8 x B200, TP=8) + 1-node Decode (8 x B200, TP=8)
   # ==============================================================================
   apiVersion: resource.k8s.io/v1
   kind: ResourceClaimTemplate
   metadata:
     name: all-mrdma
   spec:
     spec:
       devices:
         requests:
           - name: req-mrdma
             exactly:
               deviceClassName: mrdma.google.com
               allocationMode: All
   ---
   apiVersion: nvidia.com/v1alpha1
   kind: DynamoGraphDeployment
   metadata:
     name: llama3-405b-a4-disagg
     annotations:
       nvidia.com/enable-grove: "true"
   spec:
     backendFramework: sglang
     services:
       # --------------------------------------------------------------------------
       # FRONTEND: Dynamo OpenAI-Compatible Router (`python3 -m dynamo.frontend`)
       # --------------------------------------------------------------------------
       Frontend:
         componentType: frontend
         replicas: 1
         resources:
           requests:
             cpu: "8"
             memory: "32Gi"
           limits:
             cpu: "8"
             memory: "32Gi"
         extraPodSpec:
           tolerations:
             - key: "nvidia.com/gpu"
               operator: "Exists"
               effect: "NoSchedule"
           mainContainer:
             image: nvcr.io/nvidia/ai-dynamo/sglang-runtime:1.5.0
             command: ["/bin/bash", "-c"]
             args:
               - |
                 set -ex
                 exec python3 -m dynamo.frontend --http-port 8000
   
       # --------------------------------------------------------------------------
       # PREFILL WORKER: 1 x GPU A4 (`a4-highgpu-8g`, 8 x B200 GPUs, TP=8)
       # --------------------------------------------------------------------------
       Prefill:
         componentType: worker
         subComponentType: prefill
         replicas: 1
         livenessProbe:
           httpGet:
             path: /live
             port: system
           initialDelaySeconds: 600
           periodSeconds: 60
           timeoutSeconds: 30
           failureThreshold: 60
         readinessProbe:
           httpGet:
             path: /health
             port: system
           initialDelaySeconds: 120
           periodSeconds: 30
           timeoutSeconds: 30
           failureThreshold: 60
         sharedMemory:
           size: 128Gi
         resources:
           requests:
             gpu: "8"
             cpu: "128"
             memory: "1500Gi"
           limits:
             gpu: "8"
             cpu: "128"
             memory: "2000Gi"
         envs:
           - name: MODEL_PATH
             value: "nvidia/Llama-3.1-405B-Instruct-FP8"
           - name: HF_HOME
             value: "/data"
           - name: HF_TOKEN
             valueFrom:
               secretKeyRef:
                 name: hf-secret
                 key: hf_api_token
                 optional: true
           - name: DYN_SYSTEM_ENABLED
             value: "true"
           - name: DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS
             value: "[\"generate\"]"
           - name: DYN_SYSTEM_PORT
             value: "9090"
           - name: LD_LIBRARY_PATH
             value: "/usr/local/gib/lib64:/usr/local/nvidia/lib64"
           - name: GLOO_SOCKET_IFNAME
             value: "eth0"
           - name: WITH_NVIDIA_PEERMEM
             value: "0"
           # Xet CAS downloads returned unretried 429s; plain HTTP downloads back off on 429.
           - name: HF_HUB_DISABLE_XET
             value: "1"
         extraPodSpec:
           resourceClaims:
             - name: rdma
               resourceClaimTemplateName: all-mrdma
           nodeSelector:
             cloud.google.com/gke-accelerator: nvidia-b200
             cloud.google.com/gke-gpu: "true"
             kubernetes.io/arch: amd64
           tolerations:
             - key: "nvidia.com/gpu"
               operator: "Exists"
               effect: "NoSchedule"
             - key: "kubernetes.io/arch"
               operator: "Equal"
               value: "amd64"
               effect: "NoSchedule"
           mainContainer:
             image: nvcr.io/nvidia/ai-dynamo/sglang-runtime:1.5.0
             resources:
               claims:
                 - name: rdma
             startupProbe:
               httpGet:
                 path: /live
                 port: system
               initialDelaySeconds: 120
               periodSeconds: 30
               timeoutSeconds: 10
               failureThreshold: 120
             command: ["/bin/bash", "-c"]
             args:
               - |
                 set -ex
                 if [ -f /usr/local/gib/scripts/set_nccl_env.sh ]; then
                   . /usr/local/gib/scripts/set_nccl_env.sh
                 fi
   
                 # Pre-download with huggingface_hub, which backs off on HTTP 429s. Dynamo's built-in
                 # downloader (ModelExpress) aborts on the first 429 and skips fetching if weights are cached.
                 echo "=== [Dynamo Prefill Worker] Ensuring weights are cached in ${HF_HOME:-/data} ==="
                 python3 -c "import os; from huggingface_hub import snapshot_download; snapshot_download(os.environ['MODEL_PATH'], max_workers=16)"
                 # --host 0.0.0.0 is required so MooncakeKVBootstrapServer binds to 0.0.0.0:30001 instead of 127.0.0.1:30001.
                 echo "=== [Dynamo Prefill Worker] Launching dynamo.sglang ==="
                 exec python3 -m dynamo.sglang \
                   --model-path "${MODEL_PATH}" \
                   --host 0.0.0.0 \
                   --disaggregation-mode prefill \
                   --disaggregation-bootstrap-port 30001 \
                   --tensor-parallel-size 8 \
                   --kv-cache-dtype fp8_e4m3 \
                   --context-length 16384 \
                   --mem-fraction-static 0.85 \
                   --enable-metrics
             volumeMounts:
               - name: library-dir-host
                 mountPath: /usr/local/nvidia
               - name: gib
                 mountPath: /usr/local/gib
               - name: data-volume
                 mountPath: /data
           volumes:
             - name: library-dir-host
               hostPath:
                 path: /home/kubernetes/bin/nvidia
             - name: gib
               hostPath:
                 path: /home/kubernetes/bin/gib
             - name: data-volume
               emptyDir: {}
   
       # --------------------------------------------------------------------------
       # DECODE WORKER: 1 x GPU A4 (`a4-highgpu-8g`, 8 x B200 GPUs, TP=8)
       # --------------------------------------------------------------------------
       Decode:
         componentType: worker
         subComponentType: decode
         replicas: 1
         livenessProbe:
           httpGet:
             path: /live
             port: system
           initialDelaySeconds: 600
           periodSeconds: 60
           timeoutSeconds: 30
           failureThreshold: 60
         readinessProbe:
           httpGet:
             path: /health
             port: system
           initialDelaySeconds: 120
           periodSeconds: 30
           timeoutSeconds: 30
           failureThreshold: 60
         sharedMemory:
           size: 128Gi
         resources:
           requests:
             gpu: "8"
             cpu: "128"
             memory: "1500Gi"
           limits:
             gpu: "8"
             cpu: "128"
             memory: "2000Gi"
         envs:
           - name: MODEL_PATH
             value: "nvidia/Llama-3.1-405B-Instruct-FP8"
           - name: HF_HOME
             value: "/data"
           - name: HF_TOKEN
             valueFrom:
               secretKeyRef:
                 name: hf-secret
                 key: hf_api_token
                 optional: true
           - name: DYN_SYSTEM_ENABLED
             value: "true"
           - name: DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS
             value: "[\"generate\"]"
           - name: DYN_SYSTEM_PORT
             value: "9090"
           - name: LD_LIBRARY_PATH
             value: "/usr/local/gib/lib64:/usr/local/nvidia/lib64"
           - name: GLOO_SOCKET_IFNAME
             value: "eth0"
           - name: WITH_NVIDIA_PEERMEM
             value: "0"
           # Xet CAS downloads returned unretried 429s; plain HTTP downloads back off on 429.
           - name: HF_HUB_DISABLE_XET
             value: "1"
         extraPodSpec:
           resourceClaims:
             - name: rdma
               resourceClaimTemplateName: all-mrdma
           nodeSelector:
             cloud.google.com/gke-accelerator: nvidia-b200
             cloud.google.com/gke-gpu: "true"
             kubernetes.io/arch: amd64
           tolerations:
             - key: "nvidia.com/gpu"
               operator: "Exists"
               effect: "NoSchedule"
             - key: "kubernetes.io/arch"
               operator: "Equal"
               value: "amd64"
               effect: "NoSchedule"
           mainContainer:
             image: nvcr.io/nvidia/ai-dynamo/sglang-runtime:1.5.0
             resources:
               claims:
                 - name: rdma
             startupProbe:
               httpGet:
                 path: /live
                 port: system
               initialDelaySeconds: 120
               periodSeconds: 30
               timeoutSeconds: 10
               failureThreshold: 120
             command: ["/bin/bash", "-c"]
             args:
               - |
                 set -ex
                 if [ -f /usr/local/gib/scripts/set_nccl_env.sh ]; then
                   . /usr/local/gib/scripts/set_nccl_env.sh
                 fi
   
                 # Pre-download with huggingface_hub, which backs off on HTTP 429s. Dynamo's built-in
                 # downloader (ModelExpress) aborts on the first 429 and skips fetching if weights are cached.
                 echo "=== [Dynamo Decode Worker] Ensuring weights are cached in ${HF_HOME:-/data} ==="
                 python3 -c "import os; from huggingface_hub import snapshot_download; snapshot_download(os.environ['MODEL_PATH'], max_workers=16)"
                 echo "=== [Dynamo Decode Worker] Launching dynamo.sglang ==="
                 exec python3 -m dynamo.sglang \
                   --model-path "${MODEL_PATH}" \
                   --host 0.0.0.0 \
                   --disaggregation-mode decode \
                   --disaggregation-bootstrap-port 30001 \
                   --tensor-parallel-size 8 \
                   --kv-cache-dtype fp8_e4m3 \
                   --context-length 16384 \
                   --mem-fraction-static 0.85 \
                   --enable-metrics
             volumeMounts:
               - name: library-dir-host
                 mountPath: /usr/local/nvidia
               - name: gib
                 mountPath: /usr/local/gib
               - name: data-volume
                 mountPath: /data
           volumes:
             - name: library-dir-host
               hostPath:
                 path: /home/kubernetes/bin/nvidia
             - name: gib
               hostPath:
                 path: /home/kubernetes/bin/gib
             - name: data-volume
               emptyDir: {}

2. 將工作負載資訊清單套用至叢集：

   kubectl apply -f llama3.1-405b-singlehost-a4.yaml

3. 如要監控權重下載和伺服器初始化作業，請列出叢集中的 Pod，並串流預先填入和解碼 Pod 的記錄：

   kubectl get pods
   POD_NAME="YOUR_POD_NAME"
   kubectl logs -f "${POD_NAME}"

   將 YOUR_POD_NAME 替換為 kubectl get pods 輸出內容中的預先填入 Pod 或解碼 Pod 名稱。

   首次執行時，大約需要 10 到 20 分鐘，才能下載 405B FP8 模型權重並初始化 SGLang 伺服器。請等到前端、預先填入和解碼 Pod 的 kubectl get pods READY 欄顯示 1/1，再傳送要求：

   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
   

傳送推論要求

向已部署的 Llama 3.1 405B 服務傳送對話完成要求：

1. 將本機通訊埠 8000 轉送至 llama3-405b-a4-disagg-frontend 服務：

   kubectl port-forward svc/llama3-405b-a4-disagg-frontend 8000:8000

   請讓通訊埠轉送功能在這個終端機工作階段中執行，或在背景執行指令。

2. 在第二個終端機工作階段中，將對話完成要求傳送至 /v1/chat/completions 端點：

   curl http://127.0.0.1:8000/v1/chat/completions \
     -H "Content-Type: application/json" \
     -d '{
       "model": "nvidia/Llama-3.1-405B-Instruct-FP8",
       "messages": [{"role": "user", "content": "Explain KV-cache disaggregation in one sentence."}],
       "max_tokens": 64
     }' | jq .

執行基準工作負載

使用 sglang.bench_serving，從 Dynamo 前端 Pod 基準測試已部署的 Llama 3.1 405B 伺服器。將基準測試設為 1,024 個輸入詞元、1,024 個輸出詞元 (1k/1k)、2,048 個提示，並將並行數量上限設為 512：

ROUTER_POD=$(kubectl get pod -l 'nvidia.com/dynamo-component-type=frontend' -o jsonpath='{.items[0].metadata.name}')

kubectl exec "${ROUTER_POD}" -- python3 -m sglang.bench_serving \
  --backend sglang-oai-chat \
  --base-url http://localhost:8000 \
  --model nvidia/Llama-3.1-405B-Instruct-FP8 \
  --dataset-name random \
  --random-input-len 1024 \
  --random-output-len 1024 \
  --random-range-ratio 1.0 \
  --num-prompts 2048 \
  --max-concurrency 512 \
  --warmup-requests 10 | tee benchmark.log

基準測試完成後，請查看包含下列指標的輸出內容：

• Output token throughput
• Total token throughput
• TTFT
• TPOT