Reference Architecture for FP8 quantized Gemma-4-31b on TPU using Qwix library

.github/workflows/dictionary/accelerated-platforms.txt

@@ -58,7 +58,9 @@ peft
 prereqs
 pretrained
 psutil
+qtype
 qwiklabs
+qwix
 rayutil
 rtxpro
 rueth

platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/inference-perf-bench/vllm/templates/configmap-benchmark.tpl.yaml

@@ -30,7 +30,7 @@ data:
       num_workers: 20
       worker_max_concurrency: 15
       worker_max_tcp_connections: 2500
-    api: 
+    api:
       type: completion
       streaming: true
     server:
@@ -47,7 +47,7 @@ data:
       prometheus:
           scrape_interval: 15
           google_managed: true         # Whether using Google Managed Prometheus
-          filters: [] 
+          filters: []
     report:
       request_lifecycle:
         summary: true
@@ -57,7 +57,6 @@ data:
         summary: true
         per_stage: true
     storage:
-      google_cloud_storage:               
+      google_cloud_storage:
         bucket_name: ${hub_models_bucket_bench_results_name}   # Required GCS bucket
         report_file_prefix: null          # Optional filename prefix
-

platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/vllm/h100-gemma-4-26b-a4b-it/patch-vllm-args.yaml

@@ -14,4 +14,4 @@
 
 - op: add
   path: /spec/template/spec/containers/1/args/-
-  value: "--chat-template=/gcs/google/gemma-4-26b-a4b/chat_template.jinja"
+  value: "--chat-template=/gcs/google/gemma-4-26b-a4b-it/chat_template.jinja"

platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/vllm/h100-gemma-4-31b-it/patch-vllm-args.yaml

@@ -14,4 +14,4 @@
 
 - op: add
   path: /spec/template/spec/containers/1/args/-
-  value: "--chat-template=/gcs/google/gemma-4-31b/chat_template.jinja"
+  value: "--chat-template=/gcs/google/gemma-4-31b-it/chat_template.jinja"

platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/vllm/l4-gemma-4-26b-a4b-it/patch-vllm-args.yaml

@@ -14,4 +14,4 @@
 
 - op: add
   path: /spec/template/spec/containers/1/args/-
-  value: "--chat-template=/gcs/google/gemma-4-26b-a4b/chat_template.jinja"
+  value: "--chat-template=/gcs/google/gemma-4-26b-a4b-it/chat_template.jinja"

platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/vllm/l4-gemma-4-31b-it/patch-vllm-args.yaml

@@ -14,4 +14,4 @@
 
 - op: add
   path: /spec/template/spec/containers/1/args/-
-  value: "--chat-template=/gcs/google/gemma-4-31b/chat_template.jinja"
+  value: "--chat-template=/gcs/google/gemma-4-31b-it/chat_template.jinja"

platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/vllm/l4-gemma-4-e2b-it/patch-vllm-args.yaml

@@ -14,4 +14,4 @@
 
 - op: add
   path: /spec/template/spec/containers/1/args/-
-  value: "--chat-template=/gcs/google/gemma-4-e2b/chat_template.jinja"
+  value: "--chat-template=/gcs/google/gemma-4-e2b-it/chat_template.jinja"

platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/vllm/l4-gemma-4-e4b-it/patch-vllm-args.yaml

@@ -14,4 +14,4 @@
 
 - op: add
   path: /spec/template/spec/containers/1/args/-
-  value: "--chat-template=/gcs/google/gemma-4-e4b/chat_template.jinja"
+  value: "--chat-template=/gcs/google/gemma-4-e4b-it/chat_template.jinja"

platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/vllm/rtx-pro-6000-gemma-4-26b-a4b-it/patch-vllm-args.yaml

@@ -14,4 +14,4 @@
 
 - op: add
   path: /spec/template/spec/containers/1/args/-
-  value: "--chat-template=/gcs/google/gemma-4-26b-a4b/chat_template.jinja"
+  value: "--chat-template=/gcs/google/gemma-4-26b-a4b-it/chat_template.jinja"

platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-gpu/vllm/rtx-pro-6000-gemma-4-31b-it/patch-vllm-args.yaml

@@ -14,4 +14,4 @@
 
 - op: add
   path: /spec/template/spec/containers/1/args/-
-  value: "--chat-template=/gcs/google/gemma-4-31b/chat_template.jinja"
+  value: "--chat-template=/gcs/google/gemma-4-31b-it/chat_template.jinja"

platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-26b-a4b-it/patch-vllm-args.yaml

@@ -14,4 +14,4 @@
 ---
 - op: add
   path: /spec/template/spec/containers/1/args/-
-  value: "--chat-template=/gcs/google/gemma-4-26b-a4b/chat_template.jinja"
+  value: "--chat-template=/gcs/google/gemma-4-26b-a4b-it/chat_template.jinja"

platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-qwix/kustomization.yaml

@@ -0,0 +1,136 @@
+# 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: v0.26.0
+
+nameSuffix: -v6e-gemma-4-31b-qwix
+
+patches:
+  - path: patch-vllm-args.yaml
+    target:
+      kind: Deployment
+      name: vllm
+  - path: patch-env.yaml
+    target:
+      kind: Deployment
+      name: vllm
+  - 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.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
+        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
+
+resources:
+  - ../base

platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-qwix/patch-env.yaml

@@ -0,0 +1,32 @@
+# 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/env/-
+  value:
+    name: MODEL_IMPL_TYPE
+    value: "flax_nnx"
+- op: add
+  path: /spec/template/spec/containers/1/env/-
+  value:
+    name: USE_BATCHED_RPA_KERNEL
+    value: "1"
+- op: add
+  path: /spec/template/spec/containers/1/env/-
+  value:
+    name: DATA_PARALLEL_SIZE
+    valueFrom:
+      configMapKeyRef:
+        name: runtime
+        key: DATA_PARALLEL_SIZE

platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-qwix/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-2x4

platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-qwix/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: "8"
+            requests:
+              google.com/tpu: "8"

platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-qwix/patch-vllm-args.yaml

@@ -0,0 +1,38 @@
+# 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-it/chat_template.jinja"
+- op: add
+  path: /spec/template/spec/containers/1/args/-
+  value: "--data-parallel-size=$(DATA_PARALLEL_SIZE)"
+- op: add
+  path: /spec/template/spec/containers/1/args/-
+  value: "--block-size=256"
+- op: add
+  path: /spec/template/spec/containers/1/args/-
+  value: "--kv-cache-dtype=fp8"
+- op: add
+  path: /spec/template/spec/containers/1/args/-
+  value: "--limit-mm-per-prompt={\"image\": 0, \"video\": 0, \"audio\": 0}"
+- op: add
+  path: /spec/template/spec/containers/1/args/-
+  value: "--mm-processor-cache-gb=0"
+- op: add
+  path: /spec/template/spec/containers/1/args/-
+  value: '--additional_config={"quantization": {"qwix": {"rules": [{"module_path": ".*", "weight_qtype": "float8_e4m3fn", "act_qtype": "float8_e4m3fn"}]}}}'
+- op: add
+  path: /spec/template/spec/containers/1/args/-
+  value: "--load-format=runai_streamer"

platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it-qwix/runtime.env

@@ -0,0 +1,9 @@
+APP_LABEL=vllm-v6e-gemma-4-31b-it-qwix
+GPU_MEMORY_UTILIZATION=0.90
+MAX_MODEL_LEN=16384
+MODEL_ID=google/gemma-4-31b-it
+MODEL_NAME=gemma-4-31b-it-qwix
+TENSOR_PARALLEL_SIZE=4
+MODEL_IMPL_TYPE=flax_nnx
+USE_BATCHED_RPA_KERNEL=1
+DATA_PARALLEL_SIZE=2

platforms/gke/base/use-cases/inference-ref-arch/kubernetes-manifests/online-inference-tpu/vllm/v6e-gemma-4-31b-it/patch-vllm-args.yaml

@@ -14,4 +14,4 @@
 
 - op: add
   path: /spec/template/spec/containers/1/args/-
-  value: "--chat-template=/gcs/google/gemma-4-31b/chat_template.jinja"
+  value: "--chat-template=/gcs/google/gemma-4-31b-it/chat_template.jinja"