From fb8e30af401f57e3f6daa66bb8cea2e6a825dbe3 Mon Sep 17 00:00:00 2001 From: Syeda Anjum Date: Tue, 28 Jul 2026 16:29:15 +0000 Subject: [PATCH 01/11] feat(tpu): add mtp speculative decoding overlay for gemma-4 --- .../v6e-gemma-4-31b-it-mtp/kustomization.yaml | 163 ++++++++++++++++++ .../patch-downloader.yaml | 48 ++++++ .../patch-nodeselector.yaml | 24 +++ .../patch-resources.yaml | 29 ++++ .../patch-vllm-args.yaml | 23 +++ .../vllm/v6e-gemma-4-31b-it-mtp/runtime.env | 7 + 6 files changed, 294 insertions(+) create mode 100644 platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml create mode 100644 platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-downloader.yaml create mode 100644 platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-nodeselector.yaml create mode 100644 platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-resources.yaml create mode 100644 platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-vllm-args.yaml create mode 100644 platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/runtime.env diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml new file mode 100644 index 000000000..ac649e8e8 --- /dev/null +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml @@ -0,0 +1,163 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +--- +apiVersion: kustomize.config.k8s.io/v1beta1 +kind: Kustomization + +configMapGenerator: + - envs: + - runtime.env + name: runtime + namespace: replaced-by-kustomize + +images: + - name: replaced-by-kustomize + newName: vllm/vllm-tpu + newTag: gemma4 + +nameSuffix: -v6e-gemma-4-31b + +patches: + - path: patch-vllm-args.yaml + target: + kind: Deployment + name: vllm + - path: patch-downloader.yaml + - path: patch-nodeselector.yaml + - path: patch-resources.yaml + + +replacements: + - source: + fieldPath: data.APP_LABEL + kind: ConfigMap + name: runtime + targets: + - fieldPaths: + - spec.selector.matchLabels.app + - spec.template.metadata.labels.app + select: + kind: Deployment + - fieldPaths: + - spec.selector.app + select: + kind: Service + # - source: + # fieldPath: data.CONTAINER_IMAGE_URL + # kind: ConfigMap + # name: vllm + # targets: + # - fieldPaths: + # - spec.template.spec.containers.[name=inference-server].image + # select: + # kind: Deployment + - source: + fieldPath: data.INFERENCE_KUBERNETES_NAMESPACE + kind: ConfigMap + name: deployment + targets: + - fieldPaths: + - metadata.namespace + select: + kind: ConfigMap + - fieldPaths: + - metadata.namespace + select: + kind: Deployment + - fieldPaths: + - metadata.namespace + select: + kind: Service + - fieldPaths: + - metadata.namespace + select: + kind: ServiceAccount + - source: + fieldPath: data.INFERENCE_KUBERNETES_SERVICE_ACCOUNT + kind: ConfigMap + name: deployment + targets: + - fieldPaths: + - spec.template.spec.serviceAccountName + select: + kind: Deployment + - fieldPaths: + - metadata.name + select: + kind: ServiceAccount + - source: + fieldPath: data.MODEL_BUCKET_NAME + kind: ConfigMap + name: deployment + targets: + - fieldPaths: + - spec.template.spec.volumes.[name=huggingface-hub-model-bucket].csi.volumeAttributes.bucketName + - spec.template.spec.volumes.[name=huggingface-hub-drafter-model-bucket].csi.volumeAttributes.bucketName + options: + delimiter: . + index: 0 + select: + kind: Deployment + - source: + fieldPath: data.MODEL_ID + kind: ConfigMap + name: runtime + targets: + - fieldPaths: + - spec.template.spec.volumes.[name=huggingface-hub-model-bucket].csi.volumeAttributes.mountOptions + options: + delimiter: "only-dir:" + index: 1 + select: + kind: Deployment + - fieldPaths: + - spec.template.spec.containers.[name=fetch-safetensors].volumeMounts.[name=huggingface-hub-model-bucket].mountPath + - spec.template.spec.containers.[name=inference-server].volumeMounts.[name=huggingface-hub-model-bucket].mountPath + options: + delimiter: / + index: 2 + select: + kind: Deployment + - source: + fieldPath: data.MODEL_NAME + kind: ConfigMap + name: runtime + targets: + - fieldPaths: + - spec.template.metadata.labels.[ai.gke.io/model] + select: + kind: Deployment + - source: + fieldPath: data.DRAFTER_MODEL_ID + kind: ConfigMap + name: runtime + targets: + - fieldPaths: + - spec.template.spec.volumes.[name=huggingface-hub-drafter-model-bucket].csi.volumeAttributes.mountOptions + options: + delimiter: "only-dir:" + index: 1 + select: + kind: Deployment + - fieldPaths: + - spec.template.spec.containers.[name=fetch-safetensors].volumeMounts.[name=huggingface-hub-drafter-model-bucket].mountPath + - spec.template.spec.containers.[name=inference-server].volumeMounts.[name=huggingface-hub-drafter-model-bucket].mountPath + options: + delimiter: / + index: 2 + select: + kind: Deployment + +resources: + - ../base diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-downloader.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-downloader.yaml new file mode 100644 index 000000000..e35f76ed8 --- /dev/null +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-downloader.yaml @@ -0,0 +1,48 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: vllm + namespace: replaced-by-kustomize +spec: + template: + spec: + volumes: + - name: huggingface-hub-drafter-model-bucket + csi: + driver: gcsfuse.csi.storage.gke.io + volumeAttributes: + bucketName: cloud-storage-bucket-name + mountOptions: metadata-cache:ttl-secs:-1,metadata-cache:stat-cache-max-size-mb:-1,metadata-cache:type-cache-max-size-mb:-1,metadata-cache:negative-ttl-secs:0,file-cache:max-size-mb:-1,file-cache:cache-file-for-range-read:true,file-cache:enable-parallel-downloads:true,implicit-dirs,file-system:kernel-list-cache-ttl-secs:-1,only-dir:replaced-by-kustomize + skipCSIBucketAccessCheck: "true" + containers: + - name: fetch-safetensors + env: + - name: DRAFTER_MODEL_ID + valueFrom: + configMapKeyRef: + name: runtime + key: DRAFTER_MODEL_ID + volumeMounts: + - mountPath: /gcs-drafter + name: huggingface-hub-drafter-model-bucket + readOnly: true + args: + - | + echo "########### $(date) - Starting parallel-fetch-safetensors for main model: ${MODEL_ID} and drafter model: ${DRAFTER_MODEL_ID}" + ls -alR /gcs + ls -alR /gcs-drafter + find /gcs/${MODEL_ID}/*safetensors -type f | xargs -I {} -P 15 sh -c 'echo "########### $(date) - Fetching Main: {}"; dd if={} of=/dev/null' + find /gcs-drafter/${DRAFTER_MODEL_ID}/*safetensors -type f | xargs -I {} -P 15 sh -c 'echo "########### $(date) - Fetching Drafter: {}"; dd if={} of=/dev/null' + echo "########### $(date) - Finished parallel-fetch-safetensors" + sleep infinity + - name: inference-server + env: + - name: DRAFTER_MODEL_ID + valueFrom: + configMapKeyRef: + name: runtime + key: DRAFTER_MODEL_ID + volumeMounts: + - mountPath: /gcs-drafter + name: huggingface-hub-drafter-model-bucket + readOnly: true diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-nodeselector.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-nodeselector.yaml new file mode 100644 index 000000000..e0f4839a7 --- /dev/null +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-nodeselector.yaml @@ -0,0 +1,24 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: vllm + namespace: replaced-by-kustomize +spec: + template: + spec: + nodeSelector: + cloud.google.com/compute-class: tpu-v6e-2x2 diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-resources.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-resources.yaml new file mode 100644 index 000000000..117ab1b36 --- /dev/null +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-resources.yaml @@ -0,0 +1,29 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: vllm + namespace: replaced-by-kustomize +spec: + template: + spec: + containers: + - name: inference-server + resources: + limits: + google.com/tpu: "4" + requests: + google.com/tpu: "4" diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-vllm-args.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-vllm-args.yaml new file mode 100644 index 000000000..e6da77f3f --- /dev/null +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-vllm-args.yaml @@ -0,0 +1,23 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +- op: add + path: /spec/template/spec/containers/1/args/- + value: "--chat-template=/gcs/google/gemma-4-31b/chat_template.jinja" +- op: add + path: /spec/template/spec/containers/1/args/- + value: "--speculative-config={\"method\":\"mtp\",\"num_speculative_tokens\":1}" +- op: add + path: /spec/template/spec/containers/1/args/- + value: "--speculative-model=/gcs-drafter/$(DRAFTER_MODEL_ID)" diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/runtime.env b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/runtime.env new file mode 100644 index 000000000..920a24f71 --- /dev/null +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/runtime.env @@ -0,0 +1,7 @@ +APP_LABEL=vllm-v6e-gemma-4-31b-it-mtp +GPU_MEMORY_UTILIZATION=0.95 +MAX_MODEL_LEN=16384 +MODEL_ID=google/gemma-4-31b-it +DRAFTER_MODEL_ID=google/gemma-4-31b-it-assistant +MODEL_NAME=gemma-4-31b-it-mtp +TENSOR_PARALLEL_SIZE=4 From 40848eaa96acd24c85940ce829561c82a703e761 Mon Sep 17 00:00:00 2001 From: Syeda Anjum Date: Tue, 28 Jul 2026 16:34:12 +0000 Subject: [PATCH 02/11] fix(tpu): update vllm-tpu image to v0.25.0 --- .../base/_shared_config/cluster.auto.tfvars | 14 +------- .../_shared_config/llmd-shared.auto.tfvars | 15 ++------ .../patch-nodeselector.yaml | 15 +------- .../patch-resources.yaml | 34 ++++++------------- .../rtx-pro-6000-gemma-4-31b-it/runtime.env | 5 +-- .../patch-nodeselector.yaml | 14 -------- .../v6e-gemma-4-31b-it/patch-resources.yaml | 26 ++++---------- .../vllm/v6e-gemma-4-31b-it/runtime.env | 3 +- .../v6e-gemma-4-31b-it-mtp/kustomization.yaml | 2 +- 9 files changed, 26 insertions(+), 102 deletions(-) diff --git a/platforms/gke/base/_shared_config/cluster.auto.tfvars b/platforms/gke/base/_shared_config/cluster.auto.tfvars index 633bb7f1e..d9c6d7955 100644 --- a/platforms/gke/base/_shared_config/cluster.auto.tfvars +++ b/platforms/gke/base/_shared_config/cluster.auto.tfvars @@ -1,13 +1 @@ -# Copyright 2024 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. +cluster_region = "us-central1" diff --git a/platforms/gke/base/use-cases/inference-ref-arch/examples/llmd/_shared_config/llmd-shared.auto.tfvars b/platforms/gke/base/use-cases/inference-ref-arch/examples/llmd/_shared_config/llmd-shared.auto.tfvars index c37e93b74..9b0bc7d5e 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/examples/llmd/_shared_config/llmd-shared.auto.tfvars +++ b/platforms/gke/base/use-cases/inference-ref-arch/examples/llmd/_shared_config/llmd-shared.auto.tfvars @@ -1,13 +1,2 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. +llmd_model_id = "google/gemma-4-31b-it" +llmd_accelerator_type = "rtx-pro-6000" diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/patch-nodeselector.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/patch-nodeselector.yaml index 5553a4420..94b8b2200 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/patch-nodeselector.yaml +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/patch-nodeselector.yaml @@ -1,17 +1,3 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. ---- apiVersion: apps/v1 kind: Deployment metadata: @@ -22,3 +8,4 @@ spec: spec: nodeSelector: cloud.google.com/compute-class: gpu-rtx-pro-6000-96gb-x2 + cloud.google.com/gke-gpu-count: '1' diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/patch-resources.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/patch-resources.yaml index e476616df..632b17279 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/patch-resources.yaml +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/patch-resources.yaml @@ -1,17 +1,3 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. ---- apiVersion: apps/v1 kind: Deployment metadata: @@ -22,13 +8,13 @@ spec: template: spec: containers: - - name: modelserver - resources: - limits: - cpu: "10" - memory: 128G - nvidia.com/gpu: "2" - requests: - cpu: "10" - memory: 128G - nvidia.com/gpu: "2" + - name: modelserver + resources: + limits: + cpu: '10' + memory: 128G + nvidia.com/gpu: 1 + requests: + cpu: '10' + memory: 128G + nvidia.com/gpu: 1 diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/runtime.env b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/runtime.env index 1373ed97f..7eca73834 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/runtime.env +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/runtime.env @@ -1,7 +1,8 @@ APP_LABEL=vllm-rtx-pro-6000-gemma-4-31b-it GPU_MEMORY_UTILIZATION=0.95 -MAX_MODEL_LEN=30000 # reduced it due to the error ValueError: To serve at least one request with the models's max seq len (32768), (27.5 GiB KV cache is needed, which is larger than the available KV cache memory (27.19 GiB). +MAX_MODEL_LEN=19660 MODEL_ID=google/gemma-4-31b-it MODEL_NAME=Gemma-4-31B-it -TENSOR_PARALLEL_SIZE=2 +TENSOR_PARALLEL_SIZE=1 HF_HOME=/.cache/huggingface +QUANTIZATION=null diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/patch-nodeselector.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/patch-nodeselector.yaml index 893ef3266..09bd76b3d 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/patch-nodeselector.yaml +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/patch-nodeselector.yaml @@ -1,17 +1,3 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. ---- apiVersion: apps/v1 kind: Deployment metadata: diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/patch-resources.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/patch-resources.yaml index 382a1e99c..563ab722b 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/patch-resources.yaml +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/patch-resources.yaml @@ -1,17 +1,3 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. ---- apiVersion: apps/v1 kind: Deployment metadata: @@ -22,9 +8,9 @@ spec: template: spec: containers: - - name: modelserver - resources: - limits: - google.com/tpu: "4" - requests: - google.com/tpu: "4" + - name: modelserver + resources: + limits: + google.com/tpu: 4 + requests: + google.com/tpu: 4 diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/runtime.env b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/runtime.env index 36ceccddf..2f48f7b1c 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/runtime.env +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/runtime.env @@ -1,8 +1,9 @@ APP_LABEL=vllm-v6e-gemma-4-31b-it GPU_MEMORY_UTILIZATION=0.95 -MAX_MODEL_LEN=32768 +MAX_MODEL_LEN=3072 MODEL_ID=google/gemma-4-31b-it MODEL_NAME=Gemma-4-31B-it TENSOR_PARALLEL_SIZE=4 MAX_NUM_BATCHED_TOKENS=4096 HF_HOME=/.cache/huggingface +QUANTIZATION=null diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml index ac649e8e8..0b90c8296 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml @@ -24,7 +24,7 @@ configMapGenerator: images: - name: replaced-by-kustomize newName: vllm/vllm-tpu - newTag: gemma4 + newTag: v0.25.0 nameSuffix: -v6e-gemma-4-31b From dd683e50a0e5d8473f19f89ffaba768bd419a7c3 Mon Sep 17 00:00:00 2001 From: Syeda Anjum Date: Tue, 28 Jul 2026 16:45:45 +0000 Subject: [PATCH 03/11] docs: add guide for gemma-4 mtp speculative decoding on tpu --- .../online-inference-tpu/tvllm-gemma-4-mtp.md | 168 ++++++++++++++++++ 1 file changed, 168 insertions(+) create mode 100644 docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/tvllm-gemma-4-mtp.md diff --git a/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/tvllm-gemma-4-mtp.md b/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/tvllm-gemma-4-mtp.md new file mode 100644 index 000000000..bace52adf --- /dev/null +++ b/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/tvllm-gemma-4-mtp.md @@ -0,0 +1,168 @@ +# Online inference using vLLM with MTP Speculative Decoding and TPUs on Google Kubernetes Engine (GKE) + +This document implements online inference using Trillium TPUs on Google Kubernetes Engine (GKE) using vLLM with Multi-Token Prediction (MTP) Speculative Decoding enabled for the Gemma-4 model. + +Speculative decoding is a powerful optimization technique that enhances LLM inference speed without compromising output quality. Specifically, Gemma-4's MTP architecture uses an assistant (drafter) model to predict multiple tokens in parallel, which the main model then verifies in a single step. + +This example is built on top of the [GKE Inference reference architecture](/docs/platforms/gke/base/use-cases/inference-ref-arch/README.md). + +## Before you begin + +- The [GKE Inference reference implementation](/platforms/gke/base/use-cases/inference-ref-arch/terraform/README.md) is deployed and configured. + +- Get access to the models. + - For Gemma-4: + - Consent to the license on [Kaggle](https://www.kaggle.com/) using a Hugging Face account. + - **google/gemma-4-31b-it** + - **google/gemma-4-31b-it-assistant** + +- Ensure your [Hugging Face Hub **Read** access token](/platforms/gke/base/core/huggingface/initialize/README.md) has been added to Secret Manager. + +## Create and configure the Google Cloud resources + +- Deploy the online TPU resources. + + ```shell + cd ${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/terraform/online_tpu && \ + rm -rf .terraform/ terraform.tfstate* && \ + terraform init && \ + terraform plan -input=false -out=tfplan && \ + terraform apply -input=false tfplan && \ + rm tfplan + ``` + +## Download the models to Cloud Storage + +- Choose the model and drafter. + + - **Gemma 4 31B Instruction-Tuned (MTP)**: + + ```shell + export HF_MODEL_ID="google/gemma-4-31b-it" + export DRAFTER_MODEL_ID="google/gemma-4-31b-it-assistant" + ``` + +- Source the environment configuration. + + ```shell + source "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/terraform/_shared_config/scripts/set_environment_variables.sh" + ``` + +- Configure the model download job. Note that the MTP patch intercepts the download to fetch both models. + + ```shell + "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/configure_huggingface.sh" + ``` + +- Deploy the model download job. + + ```shell + kubectl apply --kustomize "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/huggingface" + ``` + +- Watch the model download job until it is complete. + + ```shell + watch --color --interval 5 --no-title \ + "kubectl --namespace=${huggingface_hub_downloader_kubernetes_namespace_name} get job/${HF_MODEL_ID_HASH}-hf-model-to-gcs | GREP_COLORS='mt=01;92' egrep --color=always -e '^' -e 'Complete' + echo '\nLogs(last 10 lines):' + kubectl --namespace=${huggingface_hub_downloader_kubernetes_namespace_name} logs job/${HF_MODEL_ID_HASH}-hf-model-to-gcs --all-containers --tail 10" + ``` + + When the job is complete, you will see the following: + + ```text + NAME STATUS COMPLETIONS DURATION AGE + XXXXXXXX-hf-model-to-gcs Complete 1/1 ### ### + ``` + + You can press `CTRL`+`c` to terminate the watch. + +- Delete the model download job. + + ```shell + kubectl delete --ignore-not-found --kustomize "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/huggingface" + ``` + +## Deploy the inference workload + +- Source the environment configuration. + + ```shell + source "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/terraform/_shared_config/scripts/set_environment_variables.sh" + ``` + +- Configure the deployment. + + ```shell + "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/configure_vllm.sh" + ``` + +- Deploy the inference workload. + + ```shell + kubectl apply --kustomize "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp" + ``` + +- Watch the deployment until it is ready. + + ```shell + watch --color --interval 5 --no-title "kubectl --namespace=${ira_online_tpu_kubernetes_namespace_name} get deployment/vllm-v6e-gemma-4-31b | GREP_COLORS='mt=01;92' egrep --color=always -e '^' -e '1/1 1 1' + echo '\nLogs(last 10 lines):' + kubectl --namespace=${ira_online_tpu_kubernetes_namespace_name} logs deployment/vllm-v6e-gemma-4-31b --all-containers --tail 10" + ``` + +- When the deployment is ready, you will see output similar to the following: + + ```text + NAME READY UP-TO-DATE AVAILABLE AGE + vllm-v6e-gemma-4-31b 1/1 1 1 ### + ``` + + You can press `CTRL`+`c` to terminate the watch. + +- Send a test request to the model. + + Start a port forward to the model service. + + ```shell + kubectl --namespace=${ira_online_tpu_kubernetes_namespace_name} port-forward service/vllm-v6e-gemma-4-31b 8000:8000 >/dev/null & \ + PF_PID=$! + ``` + + Send a test request. + + ```shell + curl http://127.0.0.1:8000/v1/chat/completions \ + --data '{ + "model": "/gcs/'${HF_MODEL_ID}'", + "messages": [ { "role": "user", "content": "Why is the sky blue?" } ] + }' \ + --header "Content-Type: application/json" \ + --request POST \ + --show-error \ + --silent | jq + ``` + + Stop the port forward. + + ```shell + kill -9 ${PF_PID} + ``` + +- Delete the workload. + + ```shell + kubectl delete --ignore-not-found --kustomize "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp" + ``` + +## Clean up + +- Destroy the online TPU resources. + + ```shell + cd ${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/terraform/online_tpu && \ + rm -rf .terraform/ terraform.tfstate* && \ + terraform init && + terraform destroy -auto-approve + ``` From 6863543f477edfebe8ec544c3fb0d0eb77536c9c Mon Sep 17 00:00:00 2001 From: Syeda Anjum Date: Tue, 28 Jul 2026 16:50:29 +0000 Subject: [PATCH 04/11] docs: add tfvars configuration steps for gemma-4 mtp --- .../online-inference-tpu/tvllm-gemma-4-mtp.md | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/tvllm-gemma-4-mtp.md b/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/tvllm-gemma-4-mtp.md index bace52adf..0032de597 100644 --- a/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/tvllm-gemma-4-mtp.md +++ b/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/tvllm-gemma-4-mtp.md @@ -10,6 +10,18 @@ This example is built on top of the [GKE Inference reference architecture](/docs - The [GKE Inference reference implementation](/platforms/gke/base/use-cases/inference-ref-arch/terraform/README.md) is deployed and configured. +- Inject these values into the appropriate tfvars files (`${ACP_REPO_DIR}/platforms/gke/base/_shared_config/platform.auto.tfvars` and `cluster.auto.tfvars`) using sed. + + ```shell + # Update platform variables + sed -i 's/^platform_name.*/platform_name = ""/g' "${ACP_REPO_DIR}/platforms/gke/base/_shared_config/platform.auto.tfvars" + grep -q "^platform_default_project_id" "${ACP_REPO_DIR}/platforms/gke/base/_shared_config/platform.auto.tfvars" || echo "platform_default_project_id = \"\"" >> "${ACP_REPO_DIR}/platforms/gke/base/_shared_config/platform.auto.tfvars" + sed -i 's/^platform_default_project_id.*/platform_default_project_id = ""/g' "${ACP_REPO_DIR}/platforms/gke/base/_shared_config/platform.auto.tfvars" + + # Update cluster variables + sed -i 's/^cluster_region.*/cluster_region = ""/g' "${ACP_REPO_DIR}/platforms/gke/base/_shared_config/cluster.auto.tfvars" + ``` + - Get access to the models. - For Gemma-4: - Consent to the license on [Kaggle](https://www.kaggle.com/) using a Hugging Face account. From 11b4a6f206b3dfedd3f49cd98f28696ca1cef7a6 Mon Sep 17 00:00:00 2001 From: Syeda Anjum Date: Tue, 28 Jul 2026 16:56:40 +0000 Subject: [PATCH 05/11] docs: fix model download steps for mtp drafter in readme --- ...llm-gemma-4-mtp.md => vllm-gemma-4-mtp.md} | 89 ++++++++++++++----- 1 file changed, 66 insertions(+), 23 deletions(-) rename docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/{tvllm-gemma-4-mtp.md => vllm-gemma-4-mtp.md} (68%) diff --git a/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/tvllm-gemma-4-mtp.md b/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/vllm-gemma-4-mtp.md similarity index 68% rename from docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/tvllm-gemma-4-mtp.md rename to docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/vllm-gemma-4-mtp.md index 0032de597..9336c3730 100644 --- a/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/tvllm-gemma-4-mtp.md +++ b/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/vllm-gemma-4-mtp.md @@ -1,16 +1,37 @@ # Online inference using vLLM with MTP Speculative Decoding and TPUs on Google Kubernetes Engine (GKE) -This document implements online inference using Trillium TPUs on Google Kubernetes Engine (GKE) using vLLM with Multi-Token Prediction (MTP) Speculative Decoding enabled for the Gemma-4 model. +This document implements online inference using Trillium TPUs on Google +Kubernetes Engine (GKE) using vLLM with Multi-Token Prediction (MTP) Speculative +Decoding enabled for the Gemma-4 model. -Speculative decoding is a powerful optimization technique that enhances LLM inference speed without compromising output quality. Specifically, Gemma-4's MTP architecture uses an assistant (drafter) model to predict multiple tokens in parallel, which the main model then verifies in a single step. +Speculative decoding is a powerful optimization technique that enhances LLM +inference speed without compromising output quality. Specifically, Gemma-4's MTP +architecture uses an assistant (drafter) model to predict multiple tokens in +parallel, which the main model then verifies in a single step. -This example is built on top of the [GKE Inference reference architecture](/docs/platforms/gke/base/use-cases/inference-ref-arch/README.md). +This example is built on top of the +[GKE Inference reference architecture](/docs/platforms/gke/base/use-cases/inference-ref-arch/README.md). ## Before you begin -- The [GKE Inference reference implementation](/platforms/gke/base/use-cases/inference-ref-arch/terraform/README.md) is deployed and configured. +- Get access to the models. + - For Gemma-4: + - Consent to the license on [Kaggle](https://www.kaggle.com/) using a + Hugging Face account. + - **google/gemma-4-31b-it** + - **google/gemma-4-31b-it-assistant** + +- Ensure your + [Hugging Face Hub **Read** access token](/platforms/gke/base/core/huggingface/initialize/README.md) + has been added to Secret Manager. -- Inject these values into the appropriate tfvars files (`${ACP_REPO_DIR}/platforms/gke/base/_shared_config/platform.auto.tfvars` and `cluster.auto.tfvars`) using sed. +## Create and configure the Google Cloud resources + +- Create a standard GKE cluster + +- Inject these values into the appropriate tfvars files + (`${ACP_REPO_DIR}/platforms/gke/base/_shared_config/platform.auto.tfvars` and + `cluster.auto.tfvars`) using sed. ```shell # Update platform variables @@ -22,15 +43,11 @@ This example is built on top of the [GKE Inference reference architecture](/docs sed -i 's/^cluster_region.*/cluster_region = ""/g' "${ACP_REPO_DIR}/platforms/gke/base/_shared_config/cluster.auto.tfvars" ``` -- Get access to the models. - - For Gemma-4: - - Consent to the license on [Kaggle](https://www.kaggle.com/) using a Hugging Face account. - - **google/gemma-4-31b-it** - - **google/gemma-4-31b-it-assistant** - -- Ensure your [Hugging Face Hub **Read** access token](/platforms/gke/base/core/huggingface/initialize/README.md) has been added to Secret Manager. + For Standard Cluster: -## Create and configure the Google Cloud resources + ```shell + ${ACP_REPO_DIR}/platforms/gke/base/core/deploy-standard.sh + ``` - Deploy the online TPU resources. @@ -45,13 +62,12 @@ This example is built on top of the [GKE Inference reference architecture](/docs ## Download the models to Cloud Storage -- Choose the model and drafter. +- Choose the main model. - **Gemma 4 31B Instruction-Tuned (MTP)**: ```shell export HF_MODEL_ID="google/gemma-4-31b-it" - export DRAFTER_MODEL_ID="google/gemma-4-31b-it-assistant" ``` - Source the environment configuration. @@ -60,19 +76,14 @@ This example is built on top of the [GKE Inference reference architecture](/docs source "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/terraform/_shared_config/scripts/set_environment_variables.sh" ``` -- Configure the model download job. Note that the MTP patch intercepts the download to fetch both models. +- Configure and deploy the main model download job. ```shell "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/configure_huggingface.sh" - ``` - -- Deploy the model download job. - - ```shell kubectl apply --kustomize "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/huggingface" ``` -- Watch the model download job until it is complete. +- Watch the main model download job until it is complete. ```shell watch --color --interval 5 --no-title \ @@ -90,7 +101,32 @@ This example is built on top of the [GKE Inference reference architecture](/docs You can press `CTRL`+`c` to terminate the watch. -- Delete the model download job. +- Delete the main model download job. + + ```shell + kubectl delete --ignore-not-found --kustomize "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/huggingface" + ``` + +- Choose the drafter model and run the download job again. + + ```shell + export HF_MODEL_ID="google/gemma-4-31b-it-assistant" + source "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/terraform/_shared_config/scripts/set_environment_variables.sh" + + "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/configure_huggingface.sh" + kubectl apply --kustomize "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/huggingface" + ``` + +- Watch the drafter model download job until it is complete. + + ```shell + watch --color --interval 5 --no-title \ + "kubectl --namespace=${huggingface_hub_downloader_kubernetes_namespace_name} get job/${HF_MODEL_ID_HASH}-hf-model-to-gcs | GREP_COLORS='mt=01;92' egrep --color=always -e '^' -e 'Complete' + echo '\nLogs(last 10 lines):' + kubectl --namespace=${huggingface_hub_downloader_kubernetes_namespace_name} logs job/${HF_MODEL_ID_HASH}-hf-model-to-gcs --all-containers --tail 10" + ``` + +- Delete the drafter model download job. ```shell kubectl delete --ignore-not-found --kustomize "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/huggingface" @@ -98,6 +134,13 @@ This example is built on top of the [GKE Inference reference architecture](/docs ## Deploy the inference workload +- Set the environment variables for both models. + + ```shell + export HF_MODEL_ID="google/gemma-4-31b-it" + export DRAFTER_MODEL_ID="google/gemma-4-31b-it-assistant" + ``` + - Source the environment configuration. ```shell From 8ee67e13913d9ed7670141fdf16ff72f748b6d66 Mon Sep 17 00:00:00 2001 From: Syeda Anjum Date: Tue, 28 Jul 2026 17:07:24 +0000 Subject: [PATCH 06/11] fix(tpu): distinguish mtp deployment name from base gemma-4 deployment --- .../online-inference-tpu/vllm-gemma-4-mtp.md | 8 ++++---- .../vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/vllm-gemma-4-mtp.md b/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/vllm-gemma-4-mtp.md index 9336c3730..95167104a 100644 --- a/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/vllm-gemma-4-mtp.md +++ b/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/vllm-gemma-4-mtp.md @@ -162,16 +162,16 @@ This example is built on top of the - Watch the deployment until it is ready. ```shell - watch --color --interval 5 --no-title "kubectl --namespace=${ira_online_tpu_kubernetes_namespace_name} get deployment/vllm-v6e-gemma-4-31b | GREP_COLORS='mt=01;92' egrep --color=always -e '^' -e '1/1 1 1' + watch --color --interval 5 --no-title "kubectl --namespace=${ira_online_tpu_kubernetes_namespace_name} get deployment/vllm-v6e-gemma-4-31b-it-mtp | GREP_COLORS='mt=01;92' egrep --color=always -e '^' -e '1/1 1 1' echo '\nLogs(last 10 lines):' - kubectl --namespace=${ira_online_tpu_kubernetes_namespace_name} logs deployment/vllm-v6e-gemma-4-31b --all-containers --tail 10" + kubectl --namespace=${ira_online_tpu_kubernetes_namespace_name} logs deployment/vllm-v6e-gemma-4-31b-it-mtp --all-containers --tail 10" ``` - When the deployment is ready, you will see output similar to the following: ```text NAME READY UP-TO-DATE AVAILABLE AGE - vllm-v6e-gemma-4-31b 1/1 1 1 ### + vllm-v6e-gemma-4-31b-it-mtp 1/1 1 1 ### ``` You can press `CTRL`+`c` to terminate the watch. @@ -181,7 +181,7 @@ This example is built on top of the Start a port forward to the model service. ```shell - kubectl --namespace=${ira_online_tpu_kubernetes_namespace_name} port-forward service/vllm-v6e-gemma-4-31b 8000:8000 >/dev/null & \ + kubectl --namespace=${ira_online_tpu_kubernetes_namespace_name} port-forward service/vllm-v6e-gemma-4-31b-it-mtp 8000:8000 >/dev/null & \ PF_PID=$! ``` diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml index 0b90c8296..f004927c1 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml @@ -26,7 +26,7 @@ images: newName: vllm/vllm-tpu newTag: v0.25.0 -nameSuffix: -v6e-gemma-4-31b +nameSuffix: -v6e-gemma-4-31b-it-mtp patches: - path: patch-vllm-args.yaml From 8868cb0bec54dd22ab10efa1abd2ed9eec91907d Mon Sep 17 00:00:00 2001 From: Syeda Anjum Date: Tue, 28 Jul 2026 17:09:29 +0000 Subject: [PATCH 07/11] chore: run prettier on docs and manifests --- .../online-inference-tpu/vllm-gemma-4-mtp.md | 3 ++- .../vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml | 1 - .../vllm/v6e-gemma-4-31b-it-mtp/patch-vllm-args.yaml | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/vllm-gemma-4-mtp.md b/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/vllm-gemma-4-mtp.md index 95167104a..ed86b64a9 100644 --- a/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/vllm-gemma-4-mtp.md +++ b/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/vllm-gemma-4-mtp.md @@ -15,6 +15,7 @@ This example is built on top of the ## Before you begin - Get access to the models. + - For Gemma-4: - Consent to the license on [Kaggle](https://www.kaggle.com/) using a Hugging Face account. @@ -112,7 +113,7 @@ This example is built on top of the ```shell export HF_MODEL_ID="google/gemma-4-31b-it-assistant" source "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/terraform/_shared_config/scripts/set_environment_variables.sh" - + "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/configure_huggingface.sh" kubectl apply --kustomize "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/huggingface" ``` diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml index f004927c1..5d7ae5431 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml @@ -37,7 +37,6 @@ patches: - path: patch-nodeselector.yaml - path: patch-resources.yaml - replacements: - source: fieldPath: data.APP_LABEL diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-vllm-args.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-vllm-args.yaml index e6da77f3f..d00392e52 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-vllm-args.yaml +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-vllm-args.yaml @@ -17,7 +17,7 @@ value: "--chat-template=/gcs/google/gemma-4-31b/chat_template.jinja" - op: add path: /spec/template/spec/containers/1/args/- - value: "--speculative-config={\"method\":\"mtp\",\"num_speculative_tokens\":1}" + value: '--speculative-config={"method":"mtp","num_speculative_tokens":1}' - op: add path: /spec/template/spec/containers/1/args/- value: "--speculative-model=/gcs-drafter/$(DRAFTER_MODEL_ID)" From dd6646d99995df73fc0353074d9005c97a317ed4 Mon Sep 17 00:00:00 2001 From: Syeda Anjum Date: Tue, 28 Jul 2026 17:17:22 +0000 Subject: [PATCH 08/11] fix: revert accidental file modifications and add missing license header to patch-downloader.yaml --- .../base/_shared_config/cluster.auto.tfvars | 14 +++++++- .../_shared_config/llmd-shared.auto.tfvars | 15 ++++++-- .../patch-nodeselector.yaml | 15 +++++++- .../patch-resources.yaml | 34 +++++++++++++------ .../rtx-pro-6000-gemma-4-31b-it/runtime.env | 5 ++- .../patch-nodeselector.yaml | 14 ++++++++ .../v6e-gemma-4-31b-it/patch-resources.yaml | 26 ++++++++++---- .../vllm/v6e-gemma-4-31b-it/runtime.env | 3 +- .../patch-downloader.yaml | 14 ++++++++ 9 files changed, 115 insertions(+), 25 deletions(-) diff --git a/platforms/gke/base/_shared_config/cluster.auto.tfvars b/platforms/gke/base/_shared_config/cluster.auto.tfvars index d9c6d7955..633bb7f1e 100644 --- a/platforms/gke/base/_shared_config/cluster.auto.tfvars +++ b/platforms/gke/base/_shared_config/cluster.auto.tfvars @@ -1 +1,13 @@ -cluster_region = "us-central1" +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/platforms/gke/base/use-cases/inference-ref-arch/examples/llmd/_shared_config/llmd-shared.auto.tfvars b/platforms/gke/base/use-cases/inference-ref-arch/examples/llmd/_shared_config/llmd-shared.auto.tfvars index 9b0bc7d5e..c37e93b74 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/examples/llmd/_shared_config/llmd-shared.auto.tfvars +++ b/platforms/gke/base/use-cases/inference-ref-arch/examples/llmd/_shared_config/llmd-shared.auto.tfvars @@ -1,2 +1,13 @@ -llmd_model_id = "google/gemma-4-31b-it" -llmd_accelerator_type = "rtx-pro-6000" +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/patch-nodeselector.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/patch-nodeselector.yaml index 94b8b2200..5553a4420 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/patch-nodeselector.yaml +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/patch-nodeselector.yaml @@ -1,3 +1,17 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +--- apiVersion: apps/v1 kind: Deployment metadata: @@ -8,4 +22,3 @@ spec: spec: nodeSelector: cloud.google.com/compute-class: gpu-rtx-pro-6000-96gb-x2 - cloud.google.com/gke-gpu-count: '1' diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/patch-resources.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/patch-resources.yaml index 632b17279..e476616df 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/patch-resources.yaml +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/patch-resources.yaml @@ -1,3 +1,17 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +--- apiVersion: apps/v1 kind: Deployment metadata: @@ -8,13 +22,13 @@ spec: template: spec: containers: - - name: modelserver - resources: - limits: - cpu: '10' - memory: 128G - nvidia.com/gpu: 1 - requests: - cpu: '10' - memory: 128G - nvidia.com/gpu: 1 + - name: modelserver + resources: + limits: + cpu: "10" + memory: 128G + nvidia.com/gpu: "2" + requests: + cpu: "10" + memory: 128G + nvidia.com/gpu: "2" diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/runtime.env b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/runtime.env index 7eca73834..1373ed97f 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/runtime.env +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/llmd-precise-prefix-cache-routing/vllm/rtx-pro-6000-gemma-4-31b-it/runtime.env @@ -1,8 +1,7 @@ APP_LABEL=vllm-rtx-pro-6000-gemma-4-31b-it GPU_MEMORY_UTILIZATION=0.95 -MAX_MODEL_LEN=19660 +MAX_MODEL_LEN=30000 # reduced it due to the error ValueError: To serve at least one request with the models's max seq len (32768), (27.5 GiB KV cache is needed, which is larger than the available KV cache memory (27.19 GiB). MODEL_ID=google/gemma-4-31b-it MODEL_NAME=Gemma-4-31B-it -TENSOR_PARALLEL_SIZE=1 +TENSOR_PARALLEL_SIZE=2 HF_HOME=/.cache/huggingface -QUANTIZATION=null diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/patch-nodeselector.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/patch-nodeselector.yaml index 09bd76b3d..893ef3266 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/patch-nodeselector.yaml +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/patch-nodeselector.yaml @@ -1,3 +1,17 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +--- apiVersion: apps/v1 kind: Deployment metadata: diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/patch-resources.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/patch-resources.yaml index 563ab722b..382a1e99c 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/patch-resources.yaml +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/patch-resources.yaml @@ -1,3 +1,17 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +--- apiVersion: apps/v1 kind: Deployment metadata: @@ -8,9 +22,9 @@ spec: template: spec: containers: - - name: modelserver - resources: - limits: - google.com/tpu: 4 - requests: - google.com/tpu: 4 + - name: modelserver + resources: + limits: + google.com/tpu: "4" + requests: + google.com/tpu: "4" diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/runtime.env b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/runtime.env index 2f48f7b1c..36ceccddf 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/runtime.env +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/llmd-optimized-baseline/vllm/v6e-gemma-4-31b-it/runtime.env @@ -1,9 +1,8 @@ APP_LABEL=vllm-v6e-gemma-4-31b-it GPU_MEMORY_UTILIZATION=0.95 -MAX_MODEL_LEN=3072 +MAX_MODEL_LEN=32768 MODEL_ID=google/gemma-4-31b-it MODEL_NAME=Gemma-4-31B-it TENSOR_PARALLEL_SIZE=4 MAX_NUM_BATCHED_TOKENS=4096 HF_HOME=/.cache/huggingface -QUANTIZATION=null diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-downloader.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-downloader.yaml index e35f76ed8..dcae6ec47 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-downloader.yaml +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-downloader.yaml @@ -1,3 +1,17 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +--- apiVersion: apps/v1 kind: Deployment metadata: From 6e59a6f00f15230d7d1ddee6e86d96f434832c86 Mon Sep 17 00:00:00 2001 From: Syeda Anjum Date: Wed, 29 Jul 2026 02:38:45 +0000 Subject: [PATCH 09/11] docs: expand MTP explanation and add benchmarking configs --- .../online-inference-tpu/vllm-gemma-4-mtp.md | 74 +++++++++++++++- .../sd-mtp/configure_benchmark.sh | 62 ++++++++++++++ .../vllm-spec-decoding/sd-mtp/job.yaml | 60 +++++++++++++ .../sd-mtp/kustomization.yaml | 85 +++++++++++++++++++ .../sd-mtp/pod-monitoring.yaml | 35 ++++++++ .../sd-mtp/set-compute-class.yaml | 25 ++++++ .../sd-mtp/templates/benchmarking.tpl.env | 8 ++ .../templates/configmap-benchmark.tpl.yaml | 62 ++++++++++++++ ...tproviderclass-huggingface-tokens.tpl.yaml | 26 ++++++ 9 files changed, 434 insertions(+), 3 deletions(-) create mode 100755 platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/configure_benchmark.sh create mode 100644 platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/job.yaml create mode 100644 platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/kustomization.yaml create mode 100644 platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/pod-monitoring.yaml create mode 100644 platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/set-compute-class.yaml create mode 100644 platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/templates/benchmarking.tpl.env create mode 100644 platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/templates/configmap-benchmark.tpl.yaml create mode 100644 platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/templates/secretproviderclass-huggingface-tokens.tpl.yaml diff --git a/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/vllm-gemma-4-mtp.md b/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/vllm-gemma-4-mtp.md index ed86b64a9..bc66a9ee8 100644 --- a/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/vllm-gemma-4-mtp.md +++ b/docs/platforms/gke/base/use-cases/inference-ref-arch/online-inference-tpu/vllm-gemma-4-mtp.md @@ -5,9 +5,41 @@ Kubernetes Engine (GKE) using vLLM with Multi-Token Prediction (MTP) Speculative Decoding enabled for the Gemma-4 model. Speculative decoding is a powerful optimization technique that enhances LLM -inference speed without compromising output quality. Specifically, Gemma-4's MTP -architecture uses an assistant (drafter) model to predict multiple tokens in -parallel, which the main model then verifies in a single step. +inference speed without compromising output quality. By drafting multiple future +tokens and verifying them in parallel, it significantly improves latency +compared to standard auto-regressive generation. + +### How MTP Differs from Other Speculative Decoding Methods + +There are several variations of speculative decoding. MTP (Multi-Token +Prediction) takes a unique approach compared to traditional methods: + +| Method | Architecture | Training Approach | Alignment & Acceptance Rate | Key Characteristics | +| :-------------------------------- | :--------------------------------------------------------------- | :---------------------------------------------------------------------- | :-------------------------------------------------------------------------------------- | :---------------------------------------------------------------------------------------------------------------------- | +| **Standard Speculative Decoding** | Separate, smaller "draft" model and large "target" model. | Models are trained completely independently. | **Lower:** Internal logic is separated, leading to frequent rejections. | Traditional approach. Simplest to implement but less efficient. | +| **Draft-Head (e.g., EAGLE)** | Additional "draft heads" attached directly to the target model. | Draft heads are trained separately _after_ base model pre-training. | **Higher:** Shares the target model's hidden states. | Avoids a separate model, but adds significant post-training complexity. | +| **MTP (Multi-Token Prediction)** | Draft modules that share the main model's internal states. | Draft modules are trained **jointly** with the base model from scratch. | **Highest:** Perfect alignment due to joint pre-training. | Gemma-4's architecture. Achieves superior efficiency without post-training add-ons. (Loaded as an "assistant" in vLLM). | +| **dSpark** | Draft-head based parallel verification architecture. | Typically requires separate training or distillation. | **High:** Uses parallel speculative verification. | _Note: Currently, dSpark is supported via the SGLang inference engine._ | +| **N-gram (Prompt Lookup)** | No extra model or head. Uses string matching against the prompt. | No training required. | **Variable:** High for repetitive/extractive tasks, very low for open-ended generation. | Completely training-free, but highly dependent on the prompt's content. | + +### Ideal Benchmarking Datasets for Speculative Decoding + +Because speculative decoding works by predicting future tokens, the achievable +speedup is directly proportional to how predictable the generated text is. For +this reason, speculative decoding architectures like MTP perform exceptionally +well on: + +- **Extractive Summarization:** Where the model frequently quotes or restates + long phrases from the source text (e.g., **CNN Daily Mail**). +- **Retrieval-Augmented Generation (RAG):** Where the context injected into the + prompt contains the literal answers the model will output. +- **Code Generation & Formatting:** Where syntax, indentation, and variable + names are highly structured and repetitive (e.g., **HumanEval** or + **ShareGPT**). + +When benchmarking MTP, using datasets like `cnn_dailymail` rather than +open-ended conversation datasets provides a much more accurate representation of +the latency improvements you can expect in enterprise use cases. This example is built on top of the [GKE Inference reference architecture](/docs/platforms/gke/base/use-cases/inference-ref-arch/README.md). @@ -212,6 +244,42 @@ This example is built on top of the kubectl delete --ignore-not-found --kustomize "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp" ``` +## Benchmark the inference workload + +- Set the variables for the benchmark. + + ```shell + export ACCELERATOR="TPU" + export APP_LABEL="vllm-v6e-gemma-4-31b-it-mtp" + export HF_MODEL_ID="google/gemma-4-31b-it" + ``` + +- Configure the benchmark manifests. + + ```shell + "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/configure_benchmark.sh" + ``` + +- Deploy the benchmark job. + + ```shell + kubectl apply --kustomize "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp" + ``` + +- Watch the benchmark job until it completes. + + ```shell + watch --color --interval 5 --no-title "kubectl --namespace=${ira_online_tpu_kubernetes_namespace_name} get job -l app=inference-perf | GREP_COLORS='mt=01;92' egrep --color=always -e '^' -e '1/1' + echo '\nLogs(last 10 lines):' + kubectl --namespace=${ira_online_tpu_kubernetes_namespace_name} logs -l app=inference-perf --tail 10" + ``` + +- Delete the benchmark workload. + + ```shell + kubectl delete --ignore-not-found --kustomize "${ACP_REPO_DIR}/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp" + ``` + ## Clean up - Destroy the online TPU resources. diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/configure_benchmark.sh b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/configure_benchmark.sh new file mode 100755 index 000000000..5fa58eb36 --- /dev/null +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/configure_benchmark.sh @@ -0,0 +1,62 @@ +#!/usr/bin/env bash + +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +set -o errexit +set -o nounset +set -o pipefail + +MY_PATH="$( + cd "$(dirname "$0")" >/dev/null 2>&1 + pwd -P +)" + +source "${MY_PATH}/../../../../terraform/_shared_config/scripts/set_environment_variables.sh" + +# Update benchmarking namespace depending on TPU or GPU selection +TARGET_FILE="${MY_PATH}/templates/benchmarking.tpl.env" +GPU_NS="${ira_online_gpu_kubernetes_namespace_name}" +TPU_NS="${ira_online_tpu_kubernetes_namespace_name}" + +# Determine the correct namespace +if [[ "$ACCELERATOR" == "GPU" ]]; then + export BENCHMARKING_KUBERNETES_NAMESPACE=$GPU_NS +elif [[ "$ACCELERATOR" == "TPU" ]]; then + export BENCHMARKING_KUBERNETES_NAMESPACE=$TPU_NS +else + echo "Error: Please specify 'GPU' or 'TPU'" + exit 1 +fi + +# Use sed to update the value +if grep -q "BENCHMARKING_KUBERNETES_NAMESPACE=" "$TARGET_FILE"; then + sed -i "s/^BENCHMARKING_KUBERNETES_NAMESPACE=.*/BENCHMARKING_KUBERNETES_NAMESPACE=$BENCHMARKING_KUBERNETES_NAMESPACE/" "$TARGET_FILE" + echo "Successfully updated $TARGET_FILE: BENCHMARKING_KUBERNETES_NAMESPACE=$BENCHMARKING_KUBERNETES_NAMESPACE" +else + echo "Variable not found" +fi + + + +envsubst < "${MY_PATH}/templates/benchmarking.tpl.env" | sponge "${MY_PATH}/benchmarking.env" + +envsubst < "${MY_PATH}/templates/configmap-benchmark.tpl.yaml" | sponge "${MY_PATH}/configmap-benchmark.yaml" + +envsubst < "${MY_PATH}/templates/secretproviderclass-huggingface-tokens.tpl.yaml" | sponge "${MY_PATH}/secretproviderclass-huggingface-tokens.yaml" + + +cd "${MY_PATH}" +SHORT_HASH=$(echo -n "${APP_LABEL}" | sha256sum | cut -c1-10) +kustomize edit set nameprefix "${SHORT_HASH}-" + diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/job.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/job.yaml new file mode 100644 index 000000000..fc7ec10ea --- /dev/null +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/job.yaml @@ -0,0 +1,60 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +--- +apiVersion: batch/v1 +kind: Job +metadata: + name: inference-perf + namespace: replaced-by-kustomize + labels: + app: inference-perf +spec: + template: + metadata: + labels: + app: inference-perf + spec: + serviceAccountName: replaced-by-kustomize + containers: + - name: inference-perf + env: + - name: HF_TOKEN_PATH + value: /var/run/secrets/huggingface.co/token + + image: quay.io/inference-perf/inference-perf:latest + imagePullPolicy: Always + command: ["inference-perf"] + args: ["--config_file", "/etc/config/config.yaml"] + volumeMounts: + - name: config-volume + mountPath: /etc/config + readOnly: true + - name: huggingface-token + mountPath: /var/run/secrets/huggingface.co + resources: + requests: + cpu: 200m + ephemeral-storage: 10Gi + memory: 10Gi + restartPolicy: Never + volumes: + - name: config-volume + configMap: + name: inference-perf-config + - csi: + driver: secrets-store-gke.csi.k8s.io + readOnly: true + volumeAttributes: + secretProviderClass: huggingface-token-read + name: huggingface-token diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/kustomization.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/kustomization.yaml new file mode 100644 index 000000000..c4b0e1dc0 --- /dev/null +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/kustomization.yaml @@ -0,0 +1,85 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +--- +apiVersion: kustomize.config.k8s.io/v1beta1 +kind: Kustomization + +configMapGenerator: +- envs: + - benchmarking.env + name: benchmark + namespace: replaced-by-kustomize + +patches: +- path: set-compute-class.yaml + +replacements: +- source: + fieldPath: data.BENCHMARKING_KUBERNETES_NAMESPACE + kind: ConfigMap + name: benchmark + targets: + - fieldPaths: + - metadata.namespace + select: + kind: ConfigMap + - fieldPaths: + - metadata.namespace + select: + kind: Job + - fieldPaths: + - metadata.namespace + select: + kind: SecretProviderClass + - fieldPaths: + - metadata.namespace + select: + kind: PodMonitoring +- source: + fieldPath: data.BENCHMARKING_KUBERNETES_SERVICE_ACCOUNT + kind: ConfigMap + name: benchmark + targets: + - fieldPaths: + - spec.template.spec.serviceAccountName + select: + kind: Job + name: inference-perf +- source: + fieldPath: metadata.name + kind: SecretProviderClass + name: huggingface-token-read + targets: + - fieldPaths: + - spec.template.spec.volumes.[name=huggingface-token].csi.volumeAttributes.secretProviderClass + select: + kind: Job + name: inference-perf +- source: + kind: ConfigMap + name: benchmark + fieldPath: data.APP_LABEL + targets: + - select: + kind: PodMonitoring + name: vllm-podmonitoring + fieldPaths: + - metadata.labels.app + - spec.selector.matchLabels.app + +resources: +- configmap-benchmark.yaml +- job.yaml +- secretproviderclass-huggingface-tokens.yaml +- pod-monitoring.yaml diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/pod-monitoring.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/pod-monitoring.yaml new file mode 100644 index 000000000..116753958 --- /dev/null +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/pod-monitoring.yaml @@ -0,0 +1,35 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +--- +apiVersion: monitoring.googleapis.com/v1 +kind: PodMonitoring +metadata: + labels: + app: APP_LABEL + name: vllm-podmonitoring + namespace: replaced-by-kustomize +spec: + endpoints: + - interval: 15s + path: "/metrics" + port: metrics + selector: + matchLabels: + app: APP_LABEL + targetLabels: + metadata: + - pod + - container + - node + diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/set-compute-class.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/set-compute-class.yaml new file mode 100644 index 000000000..37f21ac18 --- /dev/null +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/set-compute-class.yaml @@ -0,0 +1,25 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +--- +apiVersion: batch/v1 +kind: Job +metadata: + name: inference-perf + namespace: replaced-by-kustomize +spec: + template: + spec: + nodeSelector: + cloud.google.com/compute-class: model-download + diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/templates/benchmarking.tpl.env b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/templates/benchmarking.tpl.env new file mode 100644 index 000000000..42363b9a8 --- /dev/null +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/templates/benchmarking.tpl.env @@ -0,0 +1,8 @@ +BENCHMARKING_KUBERNETES_SERVICE_ACCOUNT=${ira_inference_perf_bench_kubernetes_service_account_name} +BENCHMARKING_KUBERNETES_NAMESPACE=benchmarking_ns +HUGGINGFACE_TOKEN_READ_SECRET_PROVIDER_CLASS_NAME=huggingface-token-read +RESULTS_BUCKET_NAME=${hub_models_bucket_bench_results_name} +DATASET_BUCKET_NAME=${hub_models_bucket_bench_dataset_name} +MODEL_ID=${HF_MODEL_ID} +APP_LABEL=${APP_LABEL} + diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/templates/configmap-benchmark.tpl.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/templates/configmap-benchmark.tpl.yaml new file mode 100644 index 000000000..ebeb8b793 --- /dev/null +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/templates/configmap-benchmark.tpl.yaml @@ -0,0 +1,62 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +--- +apiVersion: v1 +kind: ConfigMap +metadata: + name: inference-perf-config + namespace: ${BENCHMARKING_KUBERNETES_NAMESPACE} +data: + config.yaml: | + load: + type: constant + interval: 1.0 + sweep: + type: linear + timeout: 250 + num_stages: 7 + stage_duration: 30 + num_workers: 20 + worker_max_concurrency: 15 + worker_max_tcp_connections: 2500 + api: + type: completion + streaming: true + server: + type: vllm + model_name: /gcs/${HF_MODEL_ID} + base_url: http://${APP_LABEL}.${BENCHMARKING_KUBERNETES_NAMESPACE}.svc.cluster.local:8000 + ignore_eos: true + tokenizer: + pretrained_model_name_or_path: ${HF_MODEL_ID} + data: + type: cnn_dailymail + metrics: + type: prometheus + prometheus: + scrape_interval: 15 + google_managed: true # Whether using Google Managed Prometheus + filters: [] + report: + request_lifecycle: + summary: true + per_stage: true + per_request: true + prometheus: + summary: true + per_stage: true + storage: + google_cloud_storage: + bucket_name: ${hub_models_bucket_bench_results_name} # Required GCS bucket + report_file_prefix: null # Optional filename prefix diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/templates/secretproviderclass-huggingface-tokens.tpl.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/templates/secretproviderclass-huggingface-tokens.tpl.yaml new file mode 100644 index 000000000..c6d280263 --- /dev/null +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/templates/secretproviderclass-huggingface-tokens.tpl.yaml @@ -0,0 +1,26 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +--- +apiVersion: secrets-store.csi.x-k8s.io/v1 +kind: SecretProviderClass +metadata: + name: huggingface-token-read + namespace: replaced-by-kustomize +spec: + parameters: + secrets: | + - resourceName: "projects/${huggingface_secret_manager_project_id}/secrets/${huggingface_hub_access_token_read_secret_manager_secret_name}/versions/latest" + path: "token" + provider: gke + From 414292df89abb1e8e9dcc663d366b793c5f0cde4 Mon Sep 17 00:00:00 2001 From: Syeda Anjum Date: Wed, 29 Jul 2026 02:39:45 +0000 Subject: [PATCH 10/11] docs: add n-gram and dataset explanation, fix OOM config --- .../vllm/v6e-gemma-4-31b-it-mtp/runtime.env | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/runtime.env b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/runtime.env index 920a24f71..0a75d246c 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/runtime.env +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/runtime.env @@ -1,5 +1,5 @@ APP_LABEL=vllm-v6e-gemma-4-31b-it-mtp -GPU_MEMORY_UTILIZATION=0.95 +GPU_MEMORY_UTILIZATION=0.80 MAX_MODEL_LEN=16384 MODEL_ID=google/gemma-4-31b-it DRAFTER_MODEL_ID=google/gemma-4-31b-it-assistant From 2e79718f665055c780089e9b485d69a940115695 Mon Sep 17 00:00:00 2001 From: Syeda Anjum Date: Fri, 7 Aug 2026 17:23:40 +0000 Subject: [PATCH 11/11] feat: Update TPU Multipod (MTP) configurations for Gemma-4 --- .../container_cluster/.terraform.lock.hcl | 2 + .../sd-mtp/benchmarking.env | 8 ++ .../sd-mtp/configmap-benchmark.yaml | 62 +++++++++++++ .../sd-mtp/kustomization.yaml | 1 + ...ecretproviderclass-huggingface-tokens.yaml | 26 ++++++ .../sd-mtp/templates/benchmarking.tpl.env | 2 +- .../model-download/huggingface/downloader.env | 6 ++ .../huggingface/kustomization.yaml | 88 +++++++++---------- ...ecretproviderclass-huggingface-tokens.yaml | 37 ++++++++ .../online-inference-tpu/base/deployment.env | 3 + .../v6e-gemma-4-31b-it-mtp/kustomization.yaml | 4 + .../patch-nodeselector.yaml | 2 +- .../patch-resources.yaml | 4 +- .../patch-vllm-args.yaml | 4 +- .../vllm/v6e-gemma-4-31b-it-mtp/runtime.env | 4 +- .../inference-ref-arch.auto.tfvars | 1 + 16 files changed, 202 insertions(+), 52 deletions(-) create mode 100644 platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/benchmarking.env create mode 100644 platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/configmap-benchmark.yaml create mode 100644 platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/secretproviderclass-huggingface-tokens.yaml create mode 100644 platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/huggingface/downloader.env create mode 100644 platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/huggingface/secretproviderclass-huggingface-tokens.yaml create mode 100644 platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/base/deployment.env diff --git a/platforms/gke/base/core/container_cluster/.terraform.lock.hcl b/platforms/gke/base/core/container_cluster/.terraform.lock.hcl index ec5fd57a3..ce84dd8d4 100644 --- a/platforms/gke/base/core/container_cluster/.terraform.lock.hcl +++ b/platforms/gke/base/core/container_cluster/.terraform.lock.hcl @@ -5,6 +5,7 @@ provider "registry.terraform.io/hashicorp/google" { version = "7.28.0" constraints = "7.28.0" hashes = [ + "h1:M3DrxwI8FiHJpvq3yVX2QWZeqv5dyLt3nQ1YBm/TNXA=", "h1:TpuhjJ5O/n/AYbIu4aUNIIFukw4KOC+vQmttJqnFFQA=", "zh:078c16b9c5e9067e72070367846976b58f906d8efab6fc4fc1325661717dc9cc", "zh:08b839014b428233a3a83d15045e7559b07fc035c7f73cc1ee2694c50c4dea54", @@ -25,6 +26,7 @@ provider "registry.terraform.io/hashicorp/google-beta" { version = "7.28.0" constraints = "7.28.0" hashes = [ + "h1:CWBjU/8mG+qj4qwZFlWDVswTBjD0UoNYchIxODBX2Sg=", "h1:UEItLGPidckHle3cfjjoZL0rkkivwPMq90S0Z9tXR+w=", "zh:227765de273f3faf94778c7e55fc7884397019a93b01add1a6204a698c880e4f", "zh:40861d737bffcdf95cc0b1bc6ac915bb126e29f088dcf2fa3c6bddf856f7f7be", diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/benchmarking.env b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/benchmarking.env new file mode 100644 index 000000000..448d31eb4 --- /dev/null +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/benchmarking.env @@ -0,0 +1,8 @@ +BENCHMARKING_KUBERNETES_SERVICE_ACCOUNT=acp-tpu-mtp-inference-perf-bench +BENCHMARKING_KUBERNETES_NAMESPACE=acp-tpu-mtp-online-tpu +HUGGINGFACE_TOKEN_READ_SECRET_PROVIDER_CLASS_NAME=huggingface-token-read +RESULTS_BUCKET_NAME=acp-tpu-mtp-bench-results +DATASET_BUCKET_NAME=acp-tpu-mtp-bench-dataset +MODEL_ID=google/gemma-4-31b-it +APP_LABEL=vllm-v6e-gemma-4-31b-it-mtp + diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/configmap-benchmark.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/configmap-benchmark.yaml new file mode 100644 index 000000000..9eedb776d --- /dev/null +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/configmap-benchmark.yaml @@ -0,0 +1,62 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +--- +apiVersion: v1 +kind: ConfigMap +metadata: + name: inference-perf-config + namespace: acp-tpu-mtp-online-tpu +data: + config.yaml: | + load: + type: constant + interval: 1.0 + sweep: + type: linear + timeout: 250 + num_stages: 7 + stage_duration: 30 + num_workers: 20 + worker_max_concurrency: 15 + worker_max_tcp_connections: 2500 + api: + type: completion + streaming: true + server: + type: vllm + model_name: /gcs/google/gemma-4-31b-it + base_url: http://vllm-v6e-gemma-4-31b-it-mtp.acp-tpu-mtp-online-tpu.svc.cluster.local:8000 + ignore_eos: true + tokenizer: + pretrained_model_name_or_path: google/gemma-7b-it + data: + type: cnn_dailymail + metrics: + type: prometheus + prometheus: + scrape_interval: 15 + google_managed: true # Whether using Google Managed Prometheus + filters: [] + report: + request_lifecycle: + summary: true + per_stage: true + per_request: true + prometheus: + summary: true + per_stage: true + storage: + google_cloud_storage: + bucket_name: acp-tpu-mtp-bench-results # Required GCS bucket + report_file_prefix: null # Optional filename prefix diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/kustomization.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/kustomization.yaml index c4b0e1dc0..9a4b3dab8 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/kustomization.yaml +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/kustomization.yaml @@ -83,3 +83,4 @@ resources: - job.yaml - secretproviderclass-huggingface-tokens.yaml - pod-monitoring.yaml + diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/secretproviderclass-huggingface-tokens.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/secretproviderclass-huggingface-tokens.yaml new file mode 100644 index 000000000..c10ed00de --- /dev/null +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/secretproviderclass-huggingface-tokens.yaml @@ -0,0 +1,26 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +--- +apiVersion: secrets-store.csi.x-k8s.io/v1 +kind: SecretProviderClass +metadata: + name: huggingface-token-read + namespace: replaced-by-kustomize +spec: + parameters: + secrets: | + - resourceName: "projects/accelerated-platforms-dev/secrets/acp-tpu-mtp-huggingface-hub-access-token-read/versions/latest" + path: "token" + provider: gke + diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/templates/benchmarking.tpl.env b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/templates/benchmarking.tpl.env index 42363b9a8..4fd3e49f1 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/templates/benchmarking.tpl.env +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm-spec-decoding/sd-mtp/templates/benchmarking.tpl.env @@ -1,5 +1,5 @@ BENCHMARKING_KUBERNETES_SERVICE_ACCOUNT=${ira_inference_perf_bench_kubernetes_service_account_name} -BENCHMARKING_KUBERNETES_NAMESPACE=benchmarking_ns +BENCHMARKING_KUBERNETES_NAMESPACE=acp-tpu-mtp-online-tpu HUGGINGFACE_TOKEN_READ_SECRET_PROVIDER_CLASS_NAME=huggingface-token-read RESULTS_BUCKET_NAME=${hub_models_bucket_bench_results_name} DATASET_BUCKET_NAME=${hub_models_bucket_bench_dataset_name} diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/huggingface/downloader.env b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/huggingface/downloader.env new file mode 100644 index 000000000..3865212a4 --- /dev/null +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/huggingface/downloader.env @@ -0,0 +1,6 @@ +DOWNLOADER_KUBERNETES_NAMESPACE=acp-tpu-mtp-hf-hub-downloader +DOWNLOADER_KUBERNETES_SERVICE_ACCOUNT=acp-tpu-mtp-hf-hub-downloader +HUGGINGFACE_TOKEN_READ_SECRET_PROVIDER_CLASS_NAME=huggingface-token-read +HUGGINGFACE_TOKEN_WRITE_SECRET_PROVIDER_CLASS_NAME=huggingface-token-write +MODEL_BUCKET_NAME=accelerated-platforms-dev-acp-tpu-mtp-hf-hub-models +MODEL_ID=google/gemma-4-31b-it-assistant diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/huggingface/kustomization.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/huggingface/kustomization.yaml index 73e8fd288..a27ecb9ee 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/huggingface/kustomization.yaml +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/huggingface/kustomization.yaml @@ -11,59 +11,59 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. ---- apiVersion: kustomize.config.k8s.io/v1beta1 kind: Kustomization configMapGenerator: - - envs: - - downloader.env - name: download - namespace: replaced-by-kustomize +- envs: + - downloader.env + name: download + namespace: replaced-by-kustomize patches: - - path: set-compute-class.yaml +- path: set-compute-class.yaml replacements: - - source: - fieldPath: data.DOWNLOADER_KUBERNETES_NAMESPACE - kind: ConfigMap - name: download - targets: - - fieldPaths: - - metadata.namespace - select: - kind: ConfigMap - - fieldPaths: - - metadata.namespace - select: - kind: Job - - fieldPaths: - - metadata.namespace - select: - kind: SecretProviderClass - - source: - fieldPath: data.DOWNLOADER_KUBERNETES_SERVICE_ACCOUNT +- source: + fieldPath: data.DOWNLOADER_KUBERNETES_NAMESPACE + kind: ConfigMap + name: download + targets: + - fieldPaths: + - metadata.namespace + select: kind: ConfigMap - name: download - targets: - - fieldPaths: - - spec.template.spec.serviceAccountName - select: - kind: Job - name: hf-model-to-gcs - - source: + - fieldPaths: + - metadata.namespace + select: + kind: Job + - fieldPaths: + - metadata.namespace + select: kind: SecretProviderClass - name: huggingface-token-read - fieldPath: metadata.name - targets: - - select: - kind: Job - name: hf-model-to-gcs - fieldPaths: - - spec.template.spec.volumes.[name=huggingface-token].csi.volumeAttributes.secretProviderClass +- source: + fieldPath: data.DOWNLOADER_KUBERNETES_SERVICE_ACCOUNT + kind: ConfigMap + name: download + targets: + - fieldPaths: + - spec.template.spec.serviceAccountName + select: + kind: Job + name: hf-model-to-gcs +- source: + fieldPath: metadata.name + kind: SecretProviderClass + name: huggingface-token-read + targets: + - fieldPaths: + - spec.template.spec.volumes.[name=huggingface-token].csi.volumeAttributes.secretProviderClass + select: + kind: Job + name: hf-model-to-gcs resources: - - configmap-scripts.yaml - - job.yaml - - secretproviderclass-huggingface-tokens.yaml +- configmap-scripts.yaml +- job.yaml +- secretproviderclass-huggingface-tokens.yaml +namePrefix: 78bf0431- diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/huggingface/secretproviderclass-huggingface-tokens.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/huggingface/secretproviderclass-huggingface-tokens.yaml new file mode 100644 index 000000000..742269c86 --- /dev/null +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/model-download/huggingface/secretproviderclass-huggingface-tokens.yaml @@ -0,0 +1,37 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +--- +apiVersion: secrets-store.csi.x-k8s.io/v1 +kind: SecretProviderClass +metadata: + name: huggingface-token-read + namespace: replaced-by-kustomize +spec: + parameters: + secrets: | + - resourceName: "projects/accelerated-platforms-dev/secrets/acp-tpu-mtp-huggingface-hub-access-token-read/versions/latest" + path: "token" + provider: gke +--- +apiVersion: secrets-store.csi.x-k8s.io/v1 +kind: SecretProviderClass +metadata: + name: huggingface-token-write + namespace: replaced-by-kustomize +spec: + parameters: + secrets: | + - resourceName: "projects/accelerated-platforms-dev/secrets/acp-tpu-mtp-huggingface-hub-access-token-write/versions/latest" + path: "token" + provider: gke diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/base/deployment.env b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/base/deployment.env new file mode 100644 index 000000000..262df35bc --- /dev/null +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/base/deployment.env @@ -0,0 +1,3 @@ +INFERENCE_KUBERNETES_NAMESPACE=acp-tpu-mtp-online-tpu +INFERENCE_KUBERNETES_SERVICE_ACCOUNT=acp-tpu-mtp-online-tpu +MODEL_BUCKET_NAME=accelerated-platforms-dev-acp-tpu-mtp-hf-hub-models diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml index 5d7ae5431..ea4ad5d2e 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/kustomization.yaml @@ -160,3 +160,7 @@ replacements: resources: - ../base + +replicas: + - name: vllm + count: 3 diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-nodeselector.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-nodeselector.yaml index e0f4839a7..f984468b3 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-nodeselector.yaml +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-nodeselector.yaml @@ -21,4 +21,4 @@ spec: template: spec: nodeSelector: - cloud.google.com/compute-class: tpu-v6e-2x2 + cloud.google.com/compute-class: tpu-v6e-2x4 diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-resources.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-resources.yaml index 117ab1b36..a2f2513e0 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-resources.yaml +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-resources.yaml @@ -24,6 +24,6 @@ spec: - name: inference-server resources: limits: - google.com/tpu: "4" + google.com/tpu: "8" requests: - google.com/tpu: "4" + google.com/tpu: "8" diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-vllm-args.yaml b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-vllm-args.yaml index d00392e52..29df5ed15 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-vllm-args.yaml +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/patch-vllm-args.yaml @@ -14,10 +14,10 @@ - op: add path: /spec/template/spec/containers/1/args/- - value: "--chat-template=/gcs/google/gemma-4-31b/chat_template.jinja" + value: "--chat-template=/gcs/$(MODEL_ID)/chat_template.jinja" - op: add path: /spec/template/spec/containers/1/args/- value: '--speculative-config={"method":"mtp","num_speculative_tokens":1}' - op: add path: /spec/template/spec/containers/1/args/- - value: "--speculative-model=/gcs-drafter/$(DRAFTER_MODEL_ID)" + value: "--spec-model=/gcs-drafter/$(DRAFTER_MODEL_ID)" diff --git a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/runtime.env b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/runtime.env index 0a75d246c..4819d263e 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/runtime.env +++ b/platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-mtp/runtime.env @@ -1,7 +1,7 @@ APP_LABEL=vllm-v6e-gemma-4-31b-it-mtp -GPU_MEMORY_UTILIZATION=0.80 +GPU_MEMORY_UTILIZATION=0.45 MAX_MODEL_LEN=16384 MODEL_ID=google/gemma-4-31b-it DRAFTER_MODEL_ID=google/gemma-4-31b-it-assistant MODEL_NAME=gemma-4-31b-it-mtp -TENSOR_PARALLEL_SIZE=4 +TENSOR_PARALLEL_SIZE=8 diff --git a/platforms/gke/base/use-cases/inference-ref-arch/terraform/_shared_config/inference-ref-arch.auto.tfvars b/platforms/gke/base/use-cases/inference-ref-arch/terraform/_shared_config/inference-ref-arch.auto.tfvars index c37e93b74..d09fd96ad 100644 --- a/platforms/gke/base/use-cases/inference-ref-arch/terraform/_shared_config/inference-ref-arch.auto.tfvars +++ b/platforms/gke/base/use-cases/inference-ref-arch/terraform/_shared_config/inference-ref-arch.auto.tfvars @@ -11,3 +11,4 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +enable_tpu = true