本教程介绍如何使用 Ray Serve LLM、vLLM 和 vLLM-TorchTPU,在 Google Kubernetes Engine (GKE) 中的多主机 Cloud TPU 切片上,通过张量并行和专家并行来应用混合专家 (MoE) Gemma 4 26B-A4B-it 大语言模型 (LLM)。
本教程适用于机器学习 (ML) 工程师、平台管理员和运维人员,以及使用 Kubernetes 容器编排功能在分布式多主机 TPU 切片上处理 AI 和机器学习工作负载的数据和 AI 专家。如需详细了解我们在 Google Cloud 内容中提及的常见角色和示例任务,请参阅常见的 GKE 用户角色和任务。
在开始本教程之前,请确保您熟悉以下内容:
背景
本部分介绍本教程中使用的关键技术。
TPU
张量处理单元 (TPU) 是 Google 为机器学习工作负载定制设计的加速器。本教程同时支持 TPU Trillium (v6e) 和 Ironwood (TPU7x)。在多主机 TPU 切片中,多个 TPU 主机通过高速芯片间互连 (ICI) 网络直接通信,提供服务超出单个芯片或 TPU 主机内存容量的模型所需的聚合高带宽内存 (HBM) 和低延迟集体通信。
vLLM 和 TorchTPU
vLLM 是一个开源 LLM 推理和服务引擎。TorchTPU 是适用于 Cloud TPU 的开源 PyTorch 后端,可通过 ICI 提供设备管理和集合通信。
vLLM-TorchTPU 会向 vLLM 注册 TPU 平台插件,并通过 TorchTPU 和加速线性代数 (XLA) 编译器为 Cloud TPU 编译 PyTorch 模型前向传递。它支持使用 Ray 在多主机 TPU 切片之间实现分布式张量并行处理和专家并行处理。
在 GKE 上使用 Ray Serve 部署 LLM
Ray Serve LLM 为 /v1/chat/completions、/v1/completions 和 /v1/models 端点提供与 OpenAI 兼容的 HTTP 入口,并在 GKE 上编排分布式 vLLM 部署。启用 Ray on GKE 插件后,KubeRay Operator 及其 TPU Webhook 会将切片拓扑元数据、工作器主机名和 TorchTPU 协调变量注入到每个多主机 TPU 工作进程 Pod 中。
Ray Serve LLM 为多主机 TPU 部署提供以下功能:
- 当您在
accelerator_config字段中指定 TPU 拓扑时,拓扑感知切片放置会为每个模型副本预留一个完整的多主机 TPU 切片,并在切片中的每个 TPU 设备上调度一个 vLLM 工作器。 - 前缀感知路由会将具有共享系统提示或多轮对话记录的请求定向到在多个 TPU 切片副本之间提供服务时已在 TPU HBM 中缓存这些前缀键值对 (KV) 块的副本。
- 直接 token 流式传输可跨并发请求批量处理出站 token 流,并直接从模型副本流式传输回答,从而缩短首 token 延迟 (TTFT) 和 token 间延迟 (ITL)。
- 基于并发的副本自动扩缩会根据正在进行的请求并发量,以整个切片为增量扩缩模型副本。
目标
本教程介绍了如何执行以下操作:
- 构建一个容器映像,将 Ray Serve 叠加在预构建的 vLLM-TorchTPU 基础映像之上,然后将其推送到 Artifact Registry。
- 在 TPU Trillium (v6e) 或 Ironwood (TPU7x) 上创建 Autopilot 或 Standard 模式的 GKE 集群。
- 部署多主机 RayCluster 自定义资源,并验证 TPU 拓扑配置。
- 部署 Ray Serve LLM 应用,以在多主机 TPU 切片上通过张量并行和专家并行来应用
google/gemma-4-26B-A4B-it模型,或者使用 RayService 自定义资源以声明方式部署集群和应用。 - 向所服务的模型发送与 OpenAI 兼容的模型发现和对话补全请求。
费用
在本教程中,您将使用 Google Cloud的以下收费组件:
您可使用价格计算器根据您的预计使用量来估算费用。
完成本教程中描述的任务后,您可以通过删除所创建的资源来避免继续计费。如需了解详情,请参阅清理。
准备工作
在开始之前,请确保您已执行以下任务:
- 启用 Artifact Registry API 和 Google Kubernetes Engine API。 启用 API
- 如需使用 Google Cloud CLI 执行此任务,请安装并初始化 gcloud CLI。如果您之前安装了 gcloud CLI,请通过运行
gcloud components update命令来获取最新版本。较早版本的 gcloud CLI 可能不支持运行本文档中的命令。
- 验证您的项目在所选可用区中是否有足够的 TPU Trillium (v6e) 或 Ironwood (TPU7x) 容量配额。如需了解详情,请参阅 Cloud TPU 配额。
- 确保您拥有 Hugging Face 账号,接受
google/gemma-4-26B-A4B-it模型页面上的许可条款,并生成具有Read权限的 Hugging Face 访问令牌。 - 确保您的账号具有以下 IAM 角色:
- Kubernetes Engine Admin (
roles/container.admin) - Artifact Registry Administrator (
roles/artifactregistry.admin)
- Kubernetes Engine Admin (
准备环境
在本教程中,您将使用 Cloud Shell 或本地工作站来管理 Google Cloud上托管的资源。Cloud Shell 预安装了本教程中使用的命令行工具,包括 kubectl、docker、python3 和 gcloud CLI。
如需使用 Cloud Shell 设置您的环境,请按照以下步骤操作:
在 Google Cloud 控制台中,点击Google Cloud 控制台中的
激活 Cloud Shell 以启动 Cloud Shell 会话。此操作会在 Google Cloud 控制台的底部窗格中启动会话。
克隆
kubernetes-engine-samples代码库并切换到本教程的示例目录:git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples.git cd kubernetes-engine-samples/ai-ml/gke-ray/rayserve/llm/torchtpu创建并激活安装了 Ray CLI 的本地 Python 虚拟环境:
python3 -m venv ~/ray-env source ~/ray-env/bin/activate pip install "ray[default]==2.58.0"为您的 Google Cloud 项目和集群设置环境变量:
export PROJECT_ID=PROJECT_ID export CLUSTER_NAME=CLUSTER_NAME export REGION=REGION export ZONE=ZONE export ARTIFACT_REGISTRY=ARTIFACT_REGISTRY export HF_TOKEN=HF_TOKEN export IMAGE_URI=${REGION}-docker.pkg.dev/${PROJECT_ID}/${ARTIFACT_REGISTRY}/ray-torchtpu-vllm:latest gcloud config set project ${PROJECT_ID}替换以下内容:
PROJECT_ID:您的 Google Cloud 项目 ID。CLUSTER_NAME:GKE 集群的名称,例如ray-serve-torchtpu-cluster。REGION:Artifact Registry 仓库和 TPU 容量所在的区域,例如us-east5。ZONE:具有可用 TPU 容量的可用区,例如us-east5-b。如需了解详情,请参阅验证 GKE 中的 TPU 可用性。ARTIFACT_REGISTRY:Artifact Registry Docker 仓库的名称,例如ray-torchtpu。HF_TOKEN:具有Read权限的 Hugging Face 访问令牌。
构建自定义容器映像
构建一个容器映像,在预构建的 vLLM-TorchTPU 基础映像(包含 TorchTPU、vLLM 和 vLLM-TorchTPU 插件)之上安装 ray[default,serve] 软件包。
查看
kubernetes-engine-samples/ai-ml/gke-ray/rayserve/llm/torchtpu目录中的Dockerfile:使用 Ray Serve LLM 进行多主机 TPU 部署时,在 TPU Trillium (v6e) 上需要使用 Ray 2.56.0 或更高版本,在 Ironwood (TPU7x) 上进行每个小芯片调度时需要使用 Ray 2.58.0 或更高版本。
创建 Artifact Registry Docker 代码库,构建容器映像,并将其推送到 Artifact Registry:
gcloud artifacts repositories create ${ARTIFACT_REGISTRY} \ --repository-format=docker \ --location=${REGION} \ --description="Ray and vLLM-TorchTPU container images" gcloud auth configure-docker ${REGION}-docker.pkg.dev,us-docker.pkg.dev docker build -t ${IMAGE_URI} . docker push ${IMAGE_URI}
创建 GKE 集群
您可以在启用了代管式 Ray Operator 插件的 GKE Autopilot 或 Standard 集群中的多主机 TPU 切片上应用 LLM。
创建 GKE 集群并为
kubectl命令行工具配置凭据:Autopilot
创建启用了 Ray Operator 插件的区域级 GKE Autopilot 集群:
gcloud container clusters create-auto ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --region=${REGION} \ --cluster-version=1.36 \ --enable-ray-operator在后续步骤中部署 RayCluster 或 RayService 自定义资源时,GKE Autopilot 会自动预配与 Pod 资源请求和节点选择器匹配的 CPU 头节点和多主机 TPU 切片。
配置
kubectl命令行工具以与您的集群通信:gcloud container clusters get-credentials ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --region=${REGION}
标准
创建一个可用区级 GKE Standard 集群,其中包含一个用于托管 Ray 头部 Pod 的 CPU 节点池,并启用映像流式传输和 GKE Dataplane V2:
gcloud container clusters create ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --cluster-version=1.36 \ --machine-type=n2-standard-8 \ --num-nodes=1 \ --addons=RayOperator \ --enable-image-streaming \ --enable-dataplane-v2配置
kubectl命令行工具以与您的集群通信:gcloud container clusters get-credentials ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE}在 TPU Trillium (v6e) 上创建一个包含 16 个设备的多主机 TPU 切片节点池,该节点池具有跨 4 个主机的
4x4拓扑;或在 Ironwood (TPU7x) 上创建一个包含 16 个设备的多主机 TPU 切片节点池,该节点池具有跨 2 个主机的2x2x2拓扑:Trillium (v6e)
gcloud container node-pools create tpu-multi-host-pool \ --project=${PROJECT_ID} \ --cluster=${CLUSTER_NAME} \ --zone=${ZONE} \ --machine-type=ct6e-standard-4t \ --tpu-topology=4x4 \ --num-nodes=4Ironwood (TPU7x)
gcloud container node-pools create tpu-multi-host-pool \ --project=${PROJECT_ID} \ --cluster=${CLUSTER_NAME} \ --zone=${ZONE} \ --machine-type=tpu7x-standard-4t \ --tpu-topology=2x2x2 \ --num-nodes=2
使用您的 Hugging Face 令牌创建一个 Kubernetes Secret,以便 Ray Pod 在下载模型权重时可以向 Hugging Face 进行身份验证:
kubectl create secret generic hf-secret \ --from-literal=hf_api_token=${HF_TOKEN}
部署多主机 RayCluster 自定义资源
部署 RayCluster 自定义资源,该资源会预配一个 CPU 头 Pod、一个用于 Ray 共享内存对象存储区的内存中 /dev/shm 卷,以及一个多主机 TPU 工作进程组。
在示例代码库中查看适用于您的 TPU 版本的 RayCluster 清单:
Trillium (v6e)
Ironwood (TPU7x)
在 Ironwood (TPU7x) 上,每个物理芯片包含两个芯粒,每个四芯片主机提供八个逻辑 TPU 核心。清单会将
resources部分中的TPU字段值设置为"8",并将RAY_TPU_RESOURCE_PER_CHIP环境变量设置为"2",以便在每台主机上通告和调度 8 个逻辑TPU资源。应用清单以创建 RayCluster 自定义资源:
Trillium (v6e)
envsubst < ray-cluster.tpu-v6e-multihost.yaml | kubectl apply -f -Ironwood (TPU7x)
envsubst < ray-cluster.tpu-v7x-multihost.yaml | kubectl apply -f -等待头 Pod 和所有 TPU 工作进程 Pod 达到
Running和1/1 READY状态:kubectl get pods -l ray.io/cluster=ray-serve-torchtpu输出类似于以下内容:
Trillium (v6e)
NAME READY STATUS RESTARTS AGE ray-serve-torchtpu-head-5cj7g 1/1 Running 0 35s ray-serve-torchtpu-v6e-16-worker-6x8cg 1/1 Running 0 35s ray-serve-torchtpu-v6e-16-worker-79j5l 1/1 Running 0 35s ray-serve-torchtpu-v6e-16-worker-psxbl 1/1 Running 0 35s ray-serve-torchtpu-v6e-16-worker-rz7dv 1/1 Running 0 35sIronwood (TPU7x)
NAME READY STATUS RESTARTS AGE ray-serve-torchtpu-head-w7b2v 1/1 Running 0 2m30s ray-serve-torchtpu-tpu7x-16-worker-6zlnj 1/1 Running 0 2m30s ray-serve-torchtpu-tpu7x-16-worker-cb4zz 1/1 Running 0 2m30s验证 KubeRay TPU 网络钩子是否已将多主机切片、工作器协调和 TorchTPU 环境变量注入工作器 Pod:
kubectl get pods -l ray.io/node-type=worker \ -o=jsonpath='{range .items[0].spec.containers[0].env[*]}{.name}={.value}{"\n"}{end}' \ | grep -E "TPU_WORKER_ID|TPU_WORKER_HOSTNAMES|TPU_NAME|TORCH_TPU_TOPOLOGY|TORCH_TPU_SLICEBUILDER_ADDRESSES"输出类似于以下内容:
Trillium (v6e)
TPU_WORKER_HOSTNAMES=v6e-16-0-0.ray-serve-torchtpu-headless,v6e-16-0-1.ray-serve-torchtpu-headless,v6e-16-0-2.ray-serve-torchtpu-headless,v6e-16-0-3.ray-serve-torchtpu-headless TORCH_TPU_TOPOLOGY=4,4,1 TORCH_TPU_SLICEBUILDER_ADDRESSES=v6e-16-0-0.ray-serve-torchtpu-headless:8471,v6e-16-0-0.ray-serve-torchtpu-headless:8472,v6e-16-0-0.ray-serve-torchtpu-headless:8473,v6e-16-0-0.ray-serve-torchtpu-headless:8474,v6e-16-0-1.ray-serve-torchtpu-headless:8471,v6e-16-0-1.ray-serve-torchtpu-headless:8472,v6e-16-0-1.ray-serve-torchtpu-headless:8473,v6e-16-0-1.ray-serve-torchtpu-headless:8474,v6e-16-0-2.ray-serve-torchtpu-headless:8471,v6e-16-0-2.ray-serve-torchtpu-headless:8472,v6e-16-0-2.ray-serve-torchtpu-headless:8473,v6e-16-0-2.ray-serve-torchtpu-headless:8474,v6e-16-0-3.ray-serve-torchtpu-headless:8471,v6e-16-0-3.ray-serve-torchtpu-headless:8472,v6e-16-0-3.ray-serve-torchtpu-headless:8473,v6e-16-0-3.ray-serve-torchtpu-headless:8474 TPU_WORKER_ID=1 TPU_NAME=v6e-16-0Ironwood (TPU7x)
TPU_WORKER_HOSTNAMES=tpu7x-16-0-0.ray-serve-torchtpu-headless,tpu7x-16-0-1.ray-serve-torchtpu-headless TORCH_TPU_TOPOLOGY=2,2,2,2 TORCH_TPU_SLICEBUILDER_ADDRESSES=tpu7x-16-0-0.ray-serve-torchtpu-headless:8471,tpu7x-16-0-0.ray-serve-torchtpu-headless:8472,tpu7x-16-0-0.ray-serve-torchtpu-headless:8473,tpu7x-16-0-0.ray-serve-torchtpu-headless:8474,tpu7x-16-0-0.ray-serve-torchtpu-headless:8475,tpu7x-16-0-0.ray-serve-torchtpu-headless:8476,tpu7x-16-0-0.ray-serve-torchtpu-headless:8477,tpu7x-16-0-0.ray-serve-torchtpu-headless:8478,tpu7x-16-0-1.ray-serve-torchtpu-headless:8471,tpu7x-16-0-1.ray-serve-torchtpu-headless:8472,tpu7x-16-0-1.ray-serve-torchtpu-headless:8473,tpu7x-16-0-1.ray-serve-torchtpu-headless:8474,tpu7x-16-0-1.ray-serve-torchtpu-headless:8475,tpu7x-16-0-1.ray-serve-torchtpu-headless:8476,tpu7x-16-0-1.ray-serve-torchtpu-headless:8477,tpu7x-16-0-1.ray-serve-torchtpu-headless:8478 TPU_WORKER_ID=1 TPU_NAME=tpu7x-16-0您可以在 Ray 信息中心的集群标签页中检查 TPU 工作器节点及其自动检测到的
ray.io/tpu-*标签。将本地端口转发到端口 8265 上的 Ray 信息中心和端口 8000 上的 Ray Serve HTTP 端点:
kubectl port-forward svc/ray-serve-torchtpu-head-svc 8265:8265 8000:8000 >/dev/null 2>&1 &
创建和部署 Ray Serve LLM 应用
在 RayCluster 自定义资源上部署 Ray Serve LLM 应用,以便在多主机 TPU 切片上提供 google/gemma-4-26B-A4B-it 模型。
查看示例代码库中的
serve_vllm_torchtpu.py脚本:该脚本会使用以下设置配置
LLMConfig对象,以实现多主机 TPU 服务:accelerator_config和placement_group_config字段会在单个 TPU 切片上预留一个展示位置组,并为切片中的每个 vLLM 工作器分配一个 TPU 核心。engine_kwargs字段使用张量并行和专家并行,将模型分片到所有 16 个 TPU 核心,并启用前缀缓存。它还会停用多模态输入,以便为文本推理预留 TPU HBM;为 Gemma 4 注意力层使用统一的 KV 缓存布局;并将上下文长度上限设为 8,192 个词元,将预填充批次大小上限设为 4,096 个词元。experimental_configs字段每 10 毫秒对出站 token 流进行一次批处理,以减少并发请求的流式传输开销。runtime_env字段会为 vLLM-TorchTPU (TPU_MULTIHOST_BACKEND: "ray") 启用 Ray 多主机后端,在 TPU TensorCore(而非 SparseCore)上运行 Gemma 4 MoE 专家层,并将该设置传播到所有 16 个工作器 actor。
使用
serve run命令在 RayCluster 自定义资源上部署 Ray Serve 应用。该命令会流式传输部署日志,直到模型权重加载完毕、XLA 图在所有 16 个 TPU 设备上完成编译,并且 Serve 应用达到RUNNING状态:Trillium (v6e)
ray job submit \ --address http://localhost:8265 \ --working-dir . \ --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \ -- serve run --non-blocking serve_vllm_torchtpu:app accelerator_type=TPU-V6E topology=4x4Ironwood (TPU7x)
ray job submit \ --address http://localhost:8265 \ --working-dir . \ --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \ -- serve run --non-blocking serve_vllm_torchtpu:app accelerator_type=TPU-V7X topology=2x2x2输出类似于以下内容:
Trillium (v6e)
(ServeController pid=1637) INFO 2026-10-13 10:23:05,160 controller 1637 -- Deploying new version of Deployment(name='LLMServer:google--gemma-4-26B-A4B-it', app='default') (initial target replicas: 1). (ServeController pid=1637) INFO 2026-10-13 10:23:05,161 controller 1637 -- Deploying new version of Deployment(name='OpenAiIngress', app='default') (initial target replicas: 1). (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) INFO 10-13 10:24:06 [__init__.py:318] TPU info: node_name=v6e-16-0 | tpu_type=v6e-16 | worker_id=1 | num_chips=4 | num_cores_per_chip=1 (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) INFO 10-13 10:24:06 [__init__.py:238] Platform plugin torchtpu is activated (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) INFO 10-13 10:24:08 [core.py:112] Initializing a V1 LLM engine (v0.22.1) with config: model='google/gemma-4-26B-A4B-it', speculative_config=None, tokenizer='google/gemma-4-26B-A4B-it', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.bfloat16, max_seq_len=8192, download_dir=None, load_format=auto, tensor_parallel_size=16, pipeline_parallel_size=1, data_parallel_size=1, decode_context_parallel_size=1, dcp_comm_backend=ag_rs, disable_custom_all_reduce=True, quantization=None, quantization_config=None, enforce_eager=False, enable_return_routed_experts=False, kv_cache_dtype=auto, device_config=tpu, ... (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) (RayWorkerWrapper pid=1017, ip=10.72.14.21) INFO 10-13 10:24:30 [expert_map_manager.py:245] [EP Rank 4/16] Expert parallelism is enabled. Expert placement strategy: linear. Local/global number of experts: 8/128. Experts local to global index map: 0->32, 1->33, 2->34, 3->35, 4->36, 5->37, 6->38, 7->39. (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) (RayWorkerWrapper pid=856, ip=10.72.6.21) INFO 10-13 10:27:55 [default_loader.py:397] Loading weights took 12.76 seconds (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) (RayWorkerWrapper pid=856, ip=10.72.6.21) INFO 10-13 10:32:53 [tpu_worker.py:376] Memory statistics | total_hbm_limit_gb=31.25GiB | total_hbm_limit_cap_gb=28.75GiB | total_hbm_used_gb=3.69GiB | kv_cache_headroom_gb=5.0GiB | kv_connector_hbm_reserve_gb=0.0GiB | total_hbm_avail_gb=20.06GiB (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) INFO 10-13 10:32:54 [kv_cache_utils.py:1733] GPU KV cache size: 600,832 tokens (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) INFO 10-13 10:32:54 [kv_cache_utils.py:1734] Maximum concurrency for 8,192 tokens per request: 73.34x (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) INFO 2026-10-13 10:36:59,851 default_LLMServer:google--gemma-4-26B-A4B-it b2grwv9c -- Started vLLM engine. INFO 2026-10-13 10:37:05,997 serve 1456 -- Application 'default' is ready at http://127.0.0.1:8000/. 2026-10-13 10:37:11,447 SUCC cli.py:66 -- ------------------------------------------ 2026-10-13 10:37:11,447 SUCC cli.py:67 -- Job 'raysubmit_94DTwVzR2pvGQVwR' succeeded 2026-10-13 10:37:11,447 SUCC cli.py:68 -- ------------------------------------------Ironwood (TPU7x)
(ServeController pid=2042) INFO 2026-10-13 10:24:38,092 controller 2042 -- Deploying new version of Deployment(name='LLMServer:google--gemma-4-26B-A4B-it', app='default') (initial target replicas: 1). (ServeController pid=2042) INFO 2026-10-13 10:24:38,093 controller 2042 -- Deploying new version of Deployment(name='OpenAiIngress', app='default') (initial target replicas: 1). (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) INFO 10-13 10:24:49 [__init__.py:318] TPU info: node_name=tpu7x-16-0 | tpu_type=tpu7x-16 | worker_id=0 | num_chips=8 | num_cores_per_chip=2 (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) INFO 10-13 10:24:49 [__init__.py:238] Platform plugin torchtpu is activated (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) INFO 10-13 10:25:23 [core.py:112] Initializing a V1 LLM engine (v0.22.1) with config: model='google/gemma-4-26B-A4B-it', speculative_config=None, tokenizer='google/gemma-4-26B-A4B-it', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.bfloat16, max_seq_len=8192, download_dir=None, load_format=auto, tensor_parallel_size=16, pipeline_parallel_size=1, data_parallel_size=1, decode_context_parallel_size=1, dcp_comm_backend=ag_rs, disable_custom_all_reduce=True, quantization=None, quantization_config=None, enforce_eager=False, enable_return_routed_experts=False, kv_cache_dtype=auto, device_config=tpu, ... (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) (RayWorkerWrapper pid=602, ip=172.16.10.6) INFO 10-13 10:25:44 [expert_map_manager.py:245] [EP Rank 8/16] Expert parallelism is enabled. Expert placement strategy: linear. Local/global number of experts: 8/128. Experts local to global index map: 0->64, 1->65, 2->66, 3->67, 4->68, 5->69, 6->70, 7->71. (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) (RayWorkerWrapper pid=762) INFO 10-13 10:29:06 [default_loader.py:397] Loading weights took 11.69 seconds (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) (RayWorkerWrapper pid=604, ip=172.16.10.6) INFO 10-13 10:34:51 [tpu_worker.py:376] Memory statistics | total_hbm_limit_gb=94.75GiB | total_hbm_limit_cap_gb=87.17GiB | total_hbm_used_gb=3.74GiB | kv_cache_headroom_gb=5.0GiB | kv_connector_hbm_reserve_gb=0.0GiB | total_hbm_avail_gb=78.43GiB (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) INFO 10-13 10:34:58 [kv_cache_utils.py:1733] GPU KV cache size: 2,349,568 tokens (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) INFO 10-13 10:34:58 [kv_cache_utils.py:1734] Maximum concurrency for 8,192 tokens per request: 286.81x (ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) INFO 2026-10-13 10:40:13,586 default_LLMServer:google--gemma-4-26B-A4B-it a15k5cv3 -- Started vLLM engine. INFO 2026-10-13 10:40:18,415 serve 1953 -- Application 'default' is ready at http://127.0.0.1:8000/. 2026-10-13 10:40:24,372 SUCC cli.py:66 -- ------------------------------------------ 2026-10-13 10:40:24,372 SUCC cli.py:67 -- Job 'raysubmit_mzWQDQGDM9QSgcMN' succeeded 2026-10-13 10:40:24,372 SUCC cli.py:68 -- ------------------------------------------
查询应用的模型
Ray Serve 应用达到 RUNNING 状态后,验证已注册的模型端点,通过端口 8000 发送与 OpenAI 兼容的 /v1/chat/completions 请求,并可选择对吞吐量和延迟时间进行基准测试。
验证模型是否已在
/v1/models端点上注册:curl http://localhost:8000/v1/models输出类似于以下内容:
{"data":[{"id":"google/gemma-4-26B-A4B-it","object":"model","owned_by":"organization-owner","permission":[],"metadata":{"model_id":"google/gemma-4-26B-A4B-it","input_modality":"text","max_request_context_length":8192}}],"object":"list"}向
/v1/chat/completions端点发送对话补全提示:curl -X POST http://localhost:8000/v1/chat/completions \ -H "Content-Type: application/json" \ -d '{ "model": "google/gemma-4-26B-A4B-it", "messages": [ {"role": "system", "content": "You are a helpful AI assistant."}, {"role": "user", "content": "Explain in one sentence how tensor parallelism and expert parallelism accelerate Mixture-of-Experts LLM inference on TPUs."} ], "max_tokens": 128, "temperature": 0.0 }'Ironwood (TPU7x) 上的输出类似于以下内容:
{"id":"chatcmpl-ec30e61c-011d-49cd-989a-9ae9f59448a8","object":"chat.completion","created":1791887860,"model":"google/gemma-4-26B-A4B-it","choices":[{"index":0,"message":{"role":"assistant","content":"Tensor parallelism accelerates inference by partitioning large weight matrices across TPU cores to reduce per-device memory and computation, while expert parallelism distributes the specialized MoE layers across different devices to enable massive model scaling and parallelize the processing of different tokens through different experts.","refusal":null,"annotations":null,"audio":null,"function_call":null,"tool_calls":[],"reasoning":null},"logprobs":null,"finish_reason":"stop","stop_reason":106,"token_ids":null,"routed_experts":null}],"service_tier":null,"system_fingerprint":"vllm-0.22.1-tp16-ep-aaf76e1a","usage":{"prompt_tokens":49,"total_tokens":100,"completion_tokens":51,"prompt_tokens_details":null},"prompt_logprobs":null,"prompt_token_ids":null,"prompt_text":null,"kv_transfer_params":null}可选:使用
vllm bench serve命令对服务吞吐量和延迟时间进行基准测试。以下命令会从 Ray 头部 Pod 运行vllm bench serveHTTP 负载生成器客户端,并向端口 8000 上的 Ray Serve HTTP 代理发送 64 个并发请求:HEAD_POD=$(kubectl get pods -l ray.io/cluster=ray-serve-torchtpu,ray.io/node-type=head -o jsonpath='{.items[0].metadata.name}') kubectl exec -it ${HEAD_POD} -- vllm bench serve \ --backend openai-chat \ --endpoint /v1/chat/completions \ --model google/gemma-4-26B-A4B-it \ --dataset-name random \ --random-input-len 1024 \ --random-output-len 512 \ --max-concurrency 32 \ --num-prompts 64 \ --num-warmups 4 \ --temperature 0.0 \ --ignore-eos \ --host localhost \ --port 8000输出类似于以下内容:
Trillium (v6e)
============ Serving Benchmark Result ============ Successful requests: 64 Failed requests: 0 Maximum request concurrency: 32 Benchmark duration (s): 8.41 Total input tokens: 66392 Total generated tokens: 32768 Request throughput (req/s): 7.61 Output token throughput (tok/s): 3894.26 Peak output token throughput (tok/s): 4409.00 Peak concurrent requests: 64.00 Total token throughput (tok/s): 11784.50 ---------------Time to First Token---------------- Mean TTFT (ms): 414.57 Median TTFT (ms): 348.94 P99 TTFT (ms): 640.26 -----Time per Output Token (excl. 1st token)------ Mean TPOT (ms): 7.37 Median TPOT (ms): 7.29 P99 TPOT (ms): 8.24 ---------------Inter-token Latency---------------- Mean ITL (ms): 7.37 Median ITL (ms): 0.00 P99 ITL (ms): 25.94 ==================================================Ironwood (TPU7x)
============ Serving Benchmark Result ============ Successful requests: 64 Failed requests: 0 Maximum request concurrency: 32 Benchmark duration (s): 9.65 Total input tokens: 66392 Total generated tokens: 32768 Request throughput (req/s): 6.64 Output token throughput (tok/s): 3397.33 Peak output token throughput (tok/s): 4574.00 Peak concurrent requests: 64.00 Total token throughput (tok/s): 10280.73 ---------------Time to First Token---------------- Mean TTFT (ms): 654.11 Median TTFT (ms): 522.30 P99 TTFT (ms): 1288.23 -----Time per Output Token (excl. 1st token)------ Mean TPOT (ms): 8.13 Median TPOT (ms): 8.08 P99 TPOT (ms): 9.69 ---------------Inter-token Latency---------------- Mean ITL (ms): 8.12 Median ITL (ms): 9.31 P99 ITL (ms): 26.85 ==================================================
使用 RayService 自定义资源以声明方式部署
对于生产部署,您可以将多主机 Ray 集群和 Ray Serve LLM 应用作为单个声明式 RayService 自定义资源进行管理,而无需手动提交作业。由于 Ray Serve LLM 构建器直接在 Serve 配置中接受模型配置,因此您不需要单独的 Python 脚本。KubeRay 会预配集群、部署 Serve 应用、监控副本运行状况,并在应用准备好处理流量时在端口 8000 上创建 Kubernetes ClusterIP Service。
如果上一部分中的 RayCluster 自定义资源仍在运行,请停止端口转发并删除 RayCluster 自定义资源,以释放 TPU 切片:
pkill -f "kubectl port-forward" || true kubectl delete raycluster ray-serve-torchtpu --ignore-not-found在示例代码库中查看适用于您的 TPU 版本的 RayService 清单:
Trillium (v6e)
Ironwood (TPU7x)
应用 RayService 清单:
Trillium (v6e)
envsubst < ray-service.tpu-v6e-multihost.yaml | kubectl apply -f -Ironwood (TPU7x)
envsubst < ray-service.tpu-v7x-multihost.yaml | kubectl apply -f -等待 RayService 自定义资源完成模型权重加载和 TPU 切片上的 XLA 图编译,并达到
Running状态:kubectl get rayservice ray-serve-torchtpu kubectl get svc ray-serve-torchtpu-serve-svc输出类似于以下内容:
NAME SERVICE STATUS NUM SERVE ENDPOINTS ray-serve-torchtpu Running 3 NAME TYPE CLUSTER-IP EXTERNAL-IP PORT(S) AGE ray-serve-torchtpu-serve-svc ClusterIP 34.118.239.213 <none> 8000/TCP 29s将本地端口 8000 转发到
ray-serve-torchtpu-serve-svc服务,并使用查询所提供的模型中的curl命令查询模型:kubectl port-forward svc/ray-serve-torchtpu-serve-svc 8000:8000 >/dev/null 2>&1 &
观察模型性能
如需在 Cloud Monitoring 中监控 TPU TensorCore 利用率、HBM 用量和 Ray 集群指标,请参阅收集和查看 GKE 上的 Ray 集群的日志和指标以及在 Ray 信息中心内查看 TPU 指标。
访问 Ray 信息中心
如需在内置 Ray 信息中心内检查 Ray Serve 部署、副本 actor、节点标签和查询延迟时间的状态,请按以下步骤操作:
如果到 Ray 头服务的端口转发尚未运行,请在后台启动:
kubectl port-forward svc/ray-serve-torchtpu-head-svc 8265:8265 8000:8000 >/dev/null 2>&1 &在 Cloud Shell 中,点击
网页预览,然后选择在端口 8265 上预览,或在本地浏览器中打开
http://localhost:8265。点击服务、集群和指标标签页,查看 vLLM 部署、副本运行状况、自动检测到的
ray.io/tpu-*节点标签、请求延迟时间指标以及按节点划分的 TPU 利用率图表。
清理
为避免因本教程中使用的资源导致您的 Google Cloud 账号产生费用,请删除 RayService 或 RayCluster 自定义资源、GKE 集群和 Artifact Registry 代码库:
停止后台端口转发进程,并删除 Ray 自定义资源和 Kubernetes Secret:
pkill -f "kubectl port-forward" || true kubectl delete rayservice ray-serve-torchtpu --ignore-not-found kubectl delete raycluster ray-serve-torchtpu --ignore-not-found kubectl delete secret hf-secret --ignore-not-found删除 GKE 集群:
Autopilot
gcloud container clusters delete ${CLUSTER_NAME} \ --region=${REGION}标准
gcloud container clusters delete ${CLUSTER_NAME} \ --zone=${ZONE}删除 Artifact Registry 代码库:
gcloud artifacts repositories delete ${ARTIFACT_REGISTRY} \ --location=${REGION}
后续步骤
- 详细了解 Cloud TPU 上的 PyTorch。
- 浏览 TorchTPU 文档。
- 了解如何使用 Ray 和 TorchTPU 在多主机 TPU 上微调模型。
- 了解如何通过 Ray 使用 GKE 中的多主机 TPU 应用 Gemma 开放模型。
- 了解如何使用 KubeRay RayService 自定义资源以声明方式管理 Ray Serve 部署。
- 阅读 GKE 中的 TPU 简介。
- 浏览 Ray Serve LLM 文档。