Implantar a inferência de vLLM de TPU desagregada com llm-d no GKE

1. Introdução

Neste codelab, você vai aprender a implantar serviços de inferência desagregados e de alta performance no Google Kubernetes Engine (GKE) usando as TPUs do Google Cloud. Você vai usar o llm-d, um framework de código aberto para disponibilização distribuída de LLMs, para separar as fases de pré-preenchimento e decodificação em vários hosts de TPU, configurar o armazenamento em cache de KV compartilhado e o gateway de inferência do GKE.

Essa configuração simula um ambiente de produção para disponibilizar modelos grandes, como o Qwen3-32B, com alta capacidade de processamento e baixa latência.

Atividades deste laboratório

  • Crie uma rede VPC personalizada com MTU otimizado para tráfego de acelerador.
  • Provisione um cluster do GKE com o driver CSI do GCS Fuse e os complementos do operador do Ray.
  • Crie oito pools de nós dedicados para frações de TPU v6e (32 chips no total).
  • Configure a Identidade da carga de trabalho e as permissões para acesso ao GCS.
  • Implante llm-d para gerenciar a veiculação desagregada do modelo Qwen3-32B.
  • Verifique a implantação com um teste comparativo.

Arquitetura

![llm-d disaggregated serving architecture showing model split into 4 2x2 replicas of prefill and the same for decode]

O que é necessário

  • Ter um projeto do Google Cloud com o faturamento ativado.
  • Uma reserva do Google Cloud para recursos de TPU v6e (32 chips, ct6e-standard-4t).
  • Um token de acesso de usuário do Hugging Face para baixar pesos do modelo.
  • O Cloud Shell ou um terminal local com gcloud, kubectl e helm instalados.
  • Duração estimada:60 minutos
  • Custo estimado:este laboratório envolve recursos significativos de TPU e custará no mínimo US $60 para concluir o projeto. Siga as etapas de limpeza imediatamente após concluir os exercícios.

2. Antes de começar

Criar ou selecionar um projeto do Google Cloud

  1. No Console do Google Cloud, selecione ou crie um projeto na nuvem do Google Cloud.
  2. Verifique se o faturamento está ativado no seu projeto do Cloud.

Iniciar o Cloud Shell

  1. Clique em Ativar o Cloud Shell na parte de cima do console do Google Cloud.
  2. Verificar a autenticação:
gcloud auth list
  1. Confirme seu projeto:
gcloud config get project
  1. Defina, se necessário:
export PROJECT_ID=<YOUR_PROJECT_ID>
gcloud config set project $PROJECT_ID

Ativar APIs

Ative os serviços necessários do Google Cloud:

gcloud services enable \
    container.googleapis.com \
    compute.googleapis.com \
    iam.googleapis.com \
    cloudresourcemanager.googleapis.com

Definir variáveis de ambiente

Defina as seguintes variáveis no shell. Substitua <YOUR_ZONE> pela sua zona de TPU alocada, <YOUR_RESERVATION_NAME> pelo ID da reserva e <YOUR_HUGGING_FACE_TOKEN> pelo seu token.

export PROJECT_ID=$(gcloud config get-value project)
export ZONE="<YOUR_ZONE>" # e.g., us-east5-a
export REGION=${ZONE%-*}
export NAMESPACE=default
export CLUSTER_NAME="qwen-serving-cluster"
export GVNIC_NETWORK_PREFIX="qwen-serving"
export RESERVATION_NAME="<YOUR_RESERVATION_NAME>"
export HF_TOKEN="<YOUR_HUGGING_FACE_TOKEN>"

3. Criar rede personalizada

A veiculação desagregada exige configurações de rede específicas para lidar com o tráfego de alta largura de banda entre os nós de pré-preenchimento e decodificação.

  1. Crie a rede VPC com uma MTU grande (8896) para uma comunicação eficiente do acelerador:
    gcloud compute --project=${PROJECT_ID} \
    networks create ${GVNIC_NETWORK_PREFIX}-main \
    --subnet-mode=auto \
    --bgp-routing-mode=regional \
    --mtu=8896
  2. Crie a sub-rede para o cluster:
    gcloud compute --project=${PROJECT_ID} \
    networks subnets create ${GVNIC_NETWORK_PREFIX}-tpu \
    --network=${GVNIC_NETWORK_PREFIX}-main \
    --region=${REGION} \
    --range=10.10.0.0/18
  3. Crie uma sub-rede somente proxy necessária para a API GKE Gateway:
    gcloud compute networks subnets create ${GVNIC_NETWORK_PREFIX}-proxy \
    --purpose=REGIONAL_MANAGED_PROXY \
    --role=ACTIVE \
    --region=${REGION} \
    --network=${GVNIC_NETWORK_PREFIX}-main \
    --range=172.16.0.0/26
  4. Crie regras de firewall para permitir a comunicação interna:
    gcloud compute --project=${PROJECT_ID} firewall-rules create ${GVNIC_NETWORK_PREFIX}-allow-internal \
    --network=${GVNIC_NETWORK_PREFIX}-main \
    --allow=all \
    --source-ranges=172.16.0.0/12,10.0.0.0/8 \
    --description="Allow all internal traffic within the network."

4. Provisionar cluster do GKE

Crie um cluster padrão do GKE configurado para oferecer suporte a montagens do GCS Fuse e cargas de trabalho do operador Ray.

  1. Crie o cluster:
    gcloud container clusters create ${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --location=${REGION} \
    --release-channel=rapid \
    --machine-type=e2-standard-4 \
    --network=${GVNIC_NETWORK_PREFIX}-main \
    --subnetwork=${GVNIC_NETWORK_PREFIX}-tpu \
    --num-nodes=1 \
    --gateway-api=standard \
    --enable-managed-prometheus \
    --enable-dataplane-v2 \
    --enable-dataplane-v2-metrics \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --addons=HttpLoadBalancing,GcsFuseCsiDriver,RayOperator,HorizontalPodAutoscaling,NodeLocalDNS \
    --enable-ip-alias
  2. Recuperar credenciais do cluster:
    gcloud container clusters get-credentials ${CLUSTER_NAME} --region=${REGION}
  3. Criar um Secret do Hugging Face:
    kubectl create secret generic llm-d-hf-token \
    --from-literal=hf_api_token=${HF_TOKEN} \
    --dry-run=client -o yaml | kubectl apply -f -

5. Criar pools de nós de TPU reservados

Provisione os oito pools de nós dedicados para as frações de TPU v6e usando sua reserva.

Execute o loop a seguir para criar os oito pools de nós:

for i in {1..8}
do
  gcloud beta container node-pools create "tpu-v6e-single-$i" \
    --project=${PROJECT_ID} \
    --cluster=${CLUSTER_NAME} \
    --region=${REGION} \
    --node-locations=${ZONE} \
    --machine-type=ct6e-standard-4t \
    --tpu-topology=2x2 \
    --num-nodes=1 \
    --reservation-affinity=specific \
    --reservation=${RESERVATION_NAME} \
    --workload-metadata=GKE_METADATA &
done

Aguarde até que todos os nós sejam criados e entrem no cluster. É possível verificar o status com kubectl get nodes.

6. Implantar o serviço llm-d

Agora você vai implantar o framework llm-d para gerenciar a veiculação desagregada.

  1. Instale o Helm para implantar gráficos llm-d:
    curl -fsSL -o get_helm.sh https://raw.githubusercontent.com/helm/helm/main/scripts/get-helm-4
chmod 700 get_helm.sh
./get_helm.sh
  2. Clone o llm-d e instale as dependências necessárias:
    git clone https://github.com/llm-d/llm-d.git
# When using yq alongside Helm, you almost always want the version by Mike Farah (mikefarah/yq).  We remove the most common yq installation before reinstalling
sudo rm -rf /usr/local/bin/yq
cd llm-d
./helpers/client-setup/install-deps.sh
  3. Prepare o arquivo custom values_tpu.yaml para configurar a veiculação desagregada do cluster:
    cat <<EOF > llm-d/guides/pd-disaggregation/ms-pd/values_tpu.yaml
multinode: false

# Configure accelerator type for Google TPU
accelerator:
type: google

modelArtifacts:
uri: "hf://Qwen/Qwen3-32B"
size: 200Gi
authSecretName: "llm-d-hf-token"
name: "Qwen/Qwen3-32B"
labels:
    llm-d.ai/inference-serving: "true"
    llm-d.ai/guide: "pd-disaggregation"
    llm-d.ai/hardware-variant: "tpu"
    llm-d.ai/hardware-vendor: "google"
    llm-d.ai/model: "Qwen3-32B"

tracing:
enabled: true
otlpEndpoint: "localhost:4317"
serviceNames:
    routingProxy: "routing-proxy"
sampling:
    sampler: "always_off"
    samplerArg: "0"

routing:
servicePort: 8000
proxy:
    image: ghcr.io/llm-d/llm-d-routing-sidecar:v0.5.0
    connector: nixlv2
    secure: false

decode:
parallelism:
    tensor: 4
create: true
replicas: 4
modelCommand: custom
extraConfig:
    nodeSelector:
    cloud.google.com/gke-tpu-accelerator: "tpu-v6e-slice"
    cloud.google.com/gke-tpu-topology: "2x2"
monitoring:
    podmonitor:
    enabled: true
    portName: "vllm"
    path: "/metrics"
    interval: "30s"
containers:
    - name: "vllm"
    image: "vllm/vllm-tpu:nightly"
    command:
        - "/bin/bash"
        - "-c"
        - |
            # ROLE: kv_consumer (Receives KV cache from prefill)
            KV_CONFIG="{\"kv_connector\":\"TPUConnector\", \"kv_connector_module_path\" : \"tpu_inference.distributed.tpu_connector\", \"kv_role\":\"kv_consumer\", \"kv_ip\" : \"$POD_IP\"}"
            echo "KV_CONFIG=$KV_CONFIG"
            python3 -m vllm.entrypoints.openai.api_server \
            --model "Qwen/Qwen3-32B" \
            --port 8200 \
            --tensor-parallel-size 4 \
            --kv-transfer-config "${KV_CONFIG}" \
            --disable-uvicorn-access-log \
            --max-num-seqs 256 \
            --block-size 128 \
            --gpu-memory-utilization 0.90 \
            --max-model-len 8192
    env:
        - name: POD_IP
        valueFrom:
            fieldRef:
            fieldPath: status.podIP
        - name: TPU_SIDE_CHANNEL_PORT
        value: "9600"
        - name: TPU_KV_TRANSFER_PORT
        value: "9100"
    ports:
        - containerPort: 8200
        name: vllm
        protocol: TCP
        - containerPort: 9100
        name: tpu-kv-transfer
        protocol: TCP
        - containerPort: 9600
        name: tpu-coord
        protocol: TCP
    resources:
        limits:
        memory: 64Gi
        cpu: "16"
        google.com/tpu: 4
        requests:
        memory: 64Gi
        cpu: "16"
        google.com/tpu: 4
    mountModelVolume: true
    volumeMounts:
        - name: metrics-volume
        mountPath: /.config
        - name: shm
        mountPath: /dev/shm
        - name: torch-compile-cache
        mountPath: /.cache
    startupProbe:
        httpGet:
        path: /health
        port: vllm
        initialDelaySeconds: 15
        periodSeconds: 30
        timeoutSeconds: 5
        failureThreshold: 120
    livenessProbe:
        httpGet:
        path: /health
        port: vllm
        periodSeconds: 10
        timeoutSeconds: 5
        failureThreshold: 3
    readinessProbe:
        httpGet:
        path: /v1/models
        port: vllm
        periodSeconds: 5
        timeoutSeconds: 2
        failureThreshold: 3
volumes:
    - name: metrics-volume
    emptyDir: {}
    - name: shm
    emptyDir:
        medium: Memory
        sizeLimit: "16Gi"
    - name: torch-compile-cache
    emptyDir: {}

prefill:
parallelism:
    tensor: 4
create: true
replicas: 4
modelCommand: custom
extraConfig:
    nodeSelector:
    cloud.google.com/gke-tpu-accelerator: "tpu-v6e-slice"
    cloud.google.com/gke-tpu-topology: "2x2"
monitoring:
    podmonitor:
    enabled: true
    portName: "vllm"
    path: "/metrics"
    interval: "30s"
containers:
    - name: "vllm"
    image: "vllm/vllm-tpu:nightly"
    command:
        - "/bin/bash"
        - "-c"
        - |
            # ROLE: kv_producer (Sends KV cache to decode)
            KV_CONFIG="{\"kv_connector\":\"TPUConnector\", \"kv_connector_module_path\" : \"tpu_inference.distributed.tpu_connector\", \"kv_role\":\"kv_producer\", \"kv_ip\" : \"$POD_IP\"}"
            echo "KV_CONFIG=$KV_CONFIG"
            python3 -m vllm.entrypoints.openai.api_server \
            --model "Qwen/Qwen3-32B" \
            --port 8200 \
            --tensor-parallel-size 4 \
            --kv-transfer-config "${KV_CONFIG}" \
            --disable-uvicorn-access-log \
            --enable-chunked-prefill \
            --block-size 128 \
            --gpu-memory-utilization 0.90 \
            --max-model-len 8192
    env:
        - name: POD_IP
        valueFrom:
            fieldRef:
            fieldPath: status.podIP
        - name: TPU_SIDE_CHANNEL_PORT
        value: "9600"
        - name: TPU_KV_TRANSFER_PORT
        value: "9100"
    ports:
        - containerPort: 8200
        name: vllm
        protocol: TCP
        - containerPort: 9100
        name: tpu-kv-transfer
        protocol: TCP
        - containerPort: 9600
        name: tpu-coord
        protocol: TCP
    resources:
        limits:
        memory: 64Gi
        cpu: "16"
        google.com/tpu: 4
        requests:
        memory: 64Gi
        cpu: "16"
        google.com/tpu: 4
    mountModelVolume: true
    volumeMounts:
        - name: metrics-volume
        mountPath: /.config
        - name: shm
        mountPath: /dev/shm
        - name: torch-compile-cache
        mountPath: /.cache
    startupProbe:
        httpGet:
        path: /health
        port: vllm
        initialDelaySeconds: 15
        periodSeconds: 30
        timeoutSeconds: 5
        failureThreshold: 120
    livenessProbe:
        httpGet:
        path: /health
        port: vllm
        periodSeconds: 10
        timeoutSeconds: 5
        failureThreshold: 3
    readinessProbe:
        httpGet:
        path: /v1/models
        port: vllm
        periodSeconds: 5
        timeoutSeconds: 2
        failureThreshold: 3
volumes:
    - name: metrics-volume
    emptyDir: {}
    - name: shm
    emptyDir:
        medium: Memory
        sizeLimit: "16Gi"
    - name: torch-compile-cache
    emptyDir: {}
EOF
  4. Implante o serviço e o gateway usando o gráfico Helm do llm-d:
    cd llm-d/guides/pd-disaggregation/
helmfile apply -e gke_tpu -n $NAMESPACE
kubectl apply -f ./httproute.gke.yaml
  5. Aguarde o início dos serviços vLLM: assista aos registros do POD de decodificação e pré-preenchimento até ver "INFO: Application startup complete."
    DECODE_POD=$(kubectl get pods -l llm-d.ai/modelservice-role=decode -o jsonpath='{.items[0].metadata.name}')

# Get the first Prefill pod name
PREFILL_POD=$(kubectl get pods -l llm-d.ai/modelservice-role=prefill -o jsonpath='{.items[0].metadata.name}')

echo "Run each of these until vLLM starts successfully and then ctrl-C out"
echo "kubectl logs -f $DECODE_POD -c vllm"
echo "kubectl logs -f $PREFILL_POD -c vllm"

7. Resposta do teste de implantação

O script abaixo vai testar a conectividade com o cluster de serviço usando o GKE Inference Gateway e, em seguida, executar um teste de comparativo de mercado.

  1. Testar a conectividade e executar o comparativo de mercado:
    cat <<EOBF > ./run_benchmark.sh
#!/bin/bash

# Configuration
NAMESPACE="default"
JOB_NAME="qwen3-pd-benchmark"
MODEL_NAME="Qwen/Qwen3-32B"

echo "🔍 Discovering Gateway IP..."
GATEWAY_IP=$(kubectl get gateway -n ${NAMESPACE} -o jsonpath='{.items[0].status.addresses[0].value}')

if [ -z "$GATEWAY_IP" ]; then
    echo "❌ Error: Could not find Gateway IP. Check 'kubectl get gateway'."
    exit 1
fi

TARGET_URL="http://${GATEWAY_IP}"
echo "✅ Found Gateway at: $TARGET_URL"

echo "🗑️  Cleaning up old benchmark jobs..."
kubectl delete job $JOB_NAME --ignore-not-found=true

echo "🚀 Generating and applying Benchmark Job..."
cat <<EOF | kubectl apply -f -
apiVersion: batch/v1
kind: Job
metadata:
name: $JOB_NAME
namespace: $NAMESPACE
spec:
template:
    spec:
    containers:
    - name: llm-benchmark
        image: vllm/vllm-openai:latest
        command: ["/bin/bash", "-c"]
        args:
        - |
            # 1. Download dataset
            if [ ! -f /data/sharegpt.json ]; then
            echo "Downloading ShareGPT dataset..."
            curl -L "https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json" -o /data/sharegpt.json
            fi

            # 2. Wait for Gateway readiness
            echo "Checking connectivity to $MODEL_NAME..."
            until curl -s "$TARGET_URL/v1/models" | grep -q "$MODEL_NAME"; do
            echo "Waiting for Gateway backends to sync..."
            sleep 10
            done

            # 3. Run Benchmark
            vllm bench serve \\
            --base-url "$TARGET_URL" \\
            --model "$MODEL_NAME" \\
            --dataset-name "sharegpt" \\
            --dataset-path "/data/sharegpt.json" \\
            --request-rate 80.0 \\
            --num-prompts 2000 \\
            --tokenizer "$MODEL_NAME"
        volumeMounts:
        - name: dataset-volume
        mountPath: /data
    restartPolicy: Never
    volumes:
    - name: dataset-volume
        emptyDir: {}
EOF

echo "⏳ Job submitted. Follow logs with:"
echo "kubectl logs -f job/$JOB_NAME"
EOBF

chmod a+x ./run_benchmark.sh

./run_benchmark.sh
    Você vai ver uma saída mostrando as solicitações sendo processadas e as métricas de latência.

8. Limpar

Para evitar cobranças contínuas na sua conta do Google Cloud, exclua os recursos criados durante este codelab.

Siga estas etapas para limpar seus recursos:

# 1. Delete LeaderWorkerSet and Helm release
kubectl delete leaderworkerset qwen-simple-anywhere-cache --ignore-not-found
helm uninstall lws --namespace lws-system 2>/dev/null
kubectl delete namespace lws-system --ignore-not-found

# 2. Delete GKE Node Pools
# Note: Usually deleting the cluster deletes the node pools, 
# but explicit deletion ensures it's gone before the cluster teardown begins.
for i in {1..8}
do
	gcloud container node-pools delete "tpu-v6e-single-$i" \
	    --cluster="${CLUSTER_NAME}" \
	    --region="${REGION}" \
	    --project="${PROJECT_ID}" --quiet

done

# 3. Delete GKE Cluster
gcloud container clusters delete "${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet

echo "--- Starting IAM and Service Account Cleanup ---"

# 1. Define the full Service Account email for clarity
SA_EMAIL="tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com"

# 2. Remove Storage Bucket IAM Binding
# This removes the 'objectViewer' role from the specific bucket
gcloud storage buckets remove-iam-policy-binding gs://inf-demo-model-storage \
    --member="serviceAccount:${SA_EMAIL}" \
    --role="roles/storage.objectViewer" --quiet

# 3. Remove Workload Identity Binding
# This severs the link between the GKE KSA and the GCP SA
gcloud iam service-accounts remove-iam-policy-binding "${SA_EMAIL}" \
    --role="roles/iam.workloadIdentityUser" \
    --member="serviceAccount:${PROJECT_ID}.svc.id.goog[default/default]" --quiet

# 4. Delete the Service Account
gcloud iam service-accounts delete "${SA_EMAIL}" --project="${PROJECT_ID}" --quiet

echo "IAM cleanup complete!"

echo "--- Starting Network and Firewall Cleanup ---"

# 4. Delete Firewall Rules (Must go before the Network)
gcloud compute firewall-rules delete \
    "${GVNIC_NETWORK_PREFIX}-allow-ssh" \
    "${GVNIC_NETWORK_PREFIX}-allow-icmp" \
    "${GVNIC_NETWORK_PREFIX}-allow-internal" \
    "ray-allow-internal" \
    --project="${PROJECT_ID}" --quiet

# 5. Delete Subnets (Must go before the Network)
gcloud compute networks subnets delete "${GVNIC_NETWORK_PREFIX}-tpu" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet

gcloud compute networks subnets delete "${GVNIC_NETWORK_PREFIX}-proxy-sub" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet

gcloud compute networks subnets delete "proxy-only-subnet" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet

# 6. Finally, delete the VPC Network
gcloud compute networks delete "${GVNIC_NETWORK_PREFIX}-main" \
    --project="${PROJECT_ID}" --quiet

echo "Cleanup complete!"

9. Parabéns

Parabéns! Você implantou o Qwen3-32B em TPUs v6e desagregadas usando llm-d e GKE.

O que você aprendeu

  • Como configurar redes personalizadas para tráfego de TPU de alta velocidade.
  • Como provisionar pools de nós da TPU reservados no GKE.
  • Como implantar llm-d para separar cargas de trabalho de pré-preenchimento e decodificação.

Próximas etapas