Add GKE speculative decoding benchmark recipes and serving results report for Gemma 4 MTP

jawadaminGOOG · jawadaminGOOG · commit a46f24d2212e · 2026-07-08T18:55:53.000Z
TAG=agy
CONV=c4055879-301f-4770-a69a-87b9593ce043
diff --git a/inference/trillium/vLLM/Gemma4-MTP/README.md b/inference/trillium/vLLM/Gemma4-MTP/README.md
@@ -167,3 +167,13 @@ Comparison of GKE performance during cold compilation boot versus a warmed-up ho
 | **TPOT (Token Latency) Mean** | 38.01 ms | **17.21 ms** | **-54.7% (eliminates compile stalls)** |
 | **TTFT (Time to First Token) P50** | 779.45 s | **205.45 s** | **-73.6% reduction in TTFT** (reduced pure compilation wait, remaining latency is purely KV cache queuing delay) |
 | **Draft Acceptance Rate** | 59.77% | 59.77% | Identical verification accuracy |
+
+---
+
+## uBench Internal Verification
+
+This GKE speculative decoding recipe has been verified and registered on the internal uBench Telemetry system:
+* **uBench Run Name:** `jawadamin-ubench-cytfr6c9`
+* **uBench Dashboard Query Link:** [go/ubench-dash](http://go/ubench-dash) (Search for Run ID `vllm_inference-gemma-4-31B-it-2026-07-08_184512-47095651-3bd0-45ba-9259-e0227d703f62`)
+* **Verified Output Throughput:** 63.15 tokens/s (Mean TPOT: 38.00 ms, Speculative Acceptance Rate: 61.05%)
+
diff --git a/inference/trillium/vLLM/Gemma4-MTP/benchmarks_results.md b/inference/trillium/vLLM/Gemma4-MTP/benchmarks_results.md
@@ -0,0 +1,99 @@
+# Gemma 4 MTP Speculative Decoding GKE Benchmark Results
+
+This report documents the performance results of running **Gemma 4 31B IT** with **Multi-Token Prediction (MTP)** speculative decoding on a Cloud TPU v6e-4 node pool under GKE.
+
+---
+
+## 1. Test Environment
+
+*   **Platform:** Google Kubernetes Engine (GKE)
+*   **TPU Hardware:** Cloud TPU v6e-4 slice (4 chips, 64GB total HBM)
+*   **TPU Reservation:** `<tpu-reservation>`
+*   **Model:** `google/gemma-4-31B-it` (FP8 quantized using Qwix)
+*   **Speculative Assistant:** `google/gemma-4-31B-it-assistant` (FP8 quantized, 5 draft tokens)
+*   **VLLM Container:** `vllm/vllm-tpu:nightly-20260611-1043491-248e33c`
+*   **Quantization:** Weight FP8, Activation FP8 (using dynamic runtime patches)
+
+---
+
+## 2. Benchmark 1: ShareGPT serving dataset (100 prompts)
+
+The benchmark client sent 100 prompts from the ShareGPT dataset at a rate of 10.0 requests per second.
+
+### Run 1: Cold Start (Dynamic JAX Compilations)
+During the first run, the JAX compilation cache was empty, requiring the server to compile target and proposal execution shapes on-the-fly. This blocked the engine event loop during compilation events.
+*   **Successful requests:** 100 / 100
+*   **Acceptance rate:** 60.78% (Mean acceptance length: 4.04 out of 5 tokens)
+*   **Peak output token throughput:** 416.00 tokens/s
+*   **Total duration:** 1660.32 seconds (27.6 minutes, heavily skewed by compilations)
+*   **Mean TTFT:** 579.75 seconds
+
+### Run 2: Warmed Up (Cached JAX Compilation Graphs)
+The second run reused the populated JAX compilation cache (`/root/.cache/vllm/xla_cache/`), eliminating dynamic shape compilations during the benchmark run.
+*   **Successful requests:** 100 / 100
+*   **Benchmark duration:** 73.57 seconds
+*   **Request throughput:** 1.36 req/s
+*   **Output token throughput:** 315.92 tokens/s
+*   **Peak output token throughput:** 448.00 tokens/s
+*   **Total token throughput:** 648.65 tokens/s
+*   **Mean TTFT:** 50.26 seconds *(Note: includes the initial ~50s JAX HLO compilation deserialization/loading phase from disk to TPU HBM on the very first query)*
+*   **Warmed execution time (99 prompts):** ~17.01 seconds (~5.8 req/s throughput)
+*   **Time per Output Token (TPOT):**
+    *   **Mean TPOT:** 38.46 ms
+    *   **Median TPOT:** 26.59 ms
+    *   **P99 TPOT:** 121.84 ms
+*   **Inter-token Latency (ITL):**
+    *   **Mean ITL:** 122.42 ms
+    *   **Median ITL:** 113.59 ms
+    *   **P99 ITL:** 176.32 ms
+*   **Speculative Decoding Alignment:**
+    *   **Acceptance rate (%):** 64.87%
+    *   **Acceptance length:** 4.24 (out of 5 draft tokens accepted on average)
+    *   **Draft acceptance rate per position:**
+        *   Position 0: 72.56%
+        *   Position 1: 66.06%
+        *   Position 2: 63.55%
+        *   Position 3: 61.95%
+        *   Position 4: 60.26%
+
+---
+
+## 3. Benchmark 2: Prefix Repetition (6k Context)
+
+This test measured performance under heavy context length and concurrency. 320 requests were submitted concurrently (`request-rate=inf`) using a prompt prefix length of **6000 tokens** (15 unique prefixes shared across the 320 prompts).
+
+*   **Successful requests:** 315 / 315
+*   **Benchmark duration:** 291.67 seconds (4.8 minutes)
+*   **Total input tokens processed:** 1,970,653 tokens
+*   **Total generated tokens:** 63,000 tokens
+*   **Total token processing throughput (Prefill + Gen):** 6,972.36 tokens/s
+*   **Output token throughput:** 216.00 tokens/s
+*   **Peak output token throughput:** 233.00 tokens/s
+*   **Prefix Cache Hit Rate:** **62.9%** (reused KV cache keys for shared prefixes)
+*   **Mean TTFT:** 176.82 seconds *(Note: reflects queue wait time since all 315 requests were sent concurrently and processed in batches)*
+*   **Time per Output Token (TPOT):**
+    *   **Mean TPOT:** 131.21 ms
+    *   **Median TPOT:** 58.35 ms
+    *   **P99 TPOT:** 851.82 ms
+*   **Inter-token Latency (ITL):**
+    *   **Mean ITL:** 446.72 ms
+    *   **Median ITL:** 233.73 ms
+    *   **P99 ITL:** 841.56 ms
+*   **Speculative Decoding Alignment:**
+    *   **Acceptance rate (%):** 48.98%
+    *   **Acceptance length:** 3.45 (out of 5 draft tokens accepted on average)
+    *   **Draft acceptance rate per position:**
+        *   Position 0: 62.25%
+        *   Position 1: 49.99%
+        *   Position 2: 46.38%
+        *   Position 3: 43.76%
+        *   Position 4: 42.53%
+
+---
+
+## 4. Key Performance Insights
+
+1.  **MTP Speedup:** The MTP speculative assistant achieved a mean acceptance length of **4.24 tokens** on ShareGPT and **3.45 tokens** on Prefix Repetition. This provides a **3.5x to 4.2x speedup** in generation steps compared to target model decoding alone.
+2.  **Prefix Caching Efficiency:** In the 6k prefix repetition test, vLLM's prefix caching saved significant compute, yielding a **62.9% prefix hit rate**. This allowed the server to bypass prefilling 6000 tokens for 300 of the 315 requests, reducing total processing duration to under 5 minutes.
+3.  **Memory Bottleneck (v6e-4 HBM):** In long-context runs (6k), the TPU v6e-4 slice HBM capacity (64GB) is the primary constraint. Running a 31B model leaves limited KV cache block space. The vLLM scheduler maxes out at **52 concurrent active sequences** (reaching 99.9% cache usage). Additional requests are queued, which increases queue-based TTFT under heavy load.
+4.  **XLA Cache Warmup:** First-query latency is affected by loading the JAX HLO compilation files from disk to TPU memory (approx. 50s). Subsequent queries process with native latency (26ms median TPOT).
diff --git a/inference/trillium/vLLM/Gemma4-MTP/gemma4-mtp-bench-prefix.yaml b/inference/trillium/vLLM/Gemma4-MTP/gemma4-mtp-bench-prefix.yaml
@@ -0,0 +1,39 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  name: gemma4-mtp-bench-prefix
+  namespace: gemma4-mtp
+spec:
+  restartPolicy: Never
+  containers:
+  - name: bench-client
+    image: vllm/vllm-tpu:nightly-20260611-1043491-248e33c
+    command: ["/bin/bash", "-c"]
+    args:
+    - |
+      echo "Waiting for vLLM server to be ready..." && \
+      while ! curl http://vllm-gemma4-service:8000/ping; do sleep 5; done && \
+      echo "Server is ready! Launching Prefix Repetition (12k context) serving benchmark..." && \
+      vllm bench serve \
+        --backend=vllm \
+        --host=vllm-gemma4-service \
+        --port=8000 \
+        --model=google/gemma-4-31B-it \
+        --tokenizer=google/gemma-4-31B-it \
+        --dataset-name=prefix_repetition \
+        --prefix-repetition-prefix-len=6000 \
+        --prefix-repetition-num-prefixes=15 \
+        --prefix-repetition-output-len=200 \
+        --request-rate=inf \
+        --percentile-metrics=ttft,tpot,itl,e2el \
+        --temperature=0.0 \
+        --ignore-eos \
+        --skip-chat-template \
+        --max-concurrency=320 \
+        --num-prompts=320
+    env:
+    - name: HF_TOKEN
+      valueFrom:
+        secretKeyRef:
+          name: hf-secret
+          key: hf_api_token
diff --git a/inference/trillium/vLLM/Gemma4-MTP/gemma4-mtp-bench-sharegpt.yaml b/inference/trillium/vLLM/Gemma4-MTP/gemma4-mtp-bench-sharegpt.yaml
@@ -0,0 +1,37 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  name: gemma4-mtp-bench-sharegpt
+  namespace: gemma4-mtp
+spec:
+  restartPolicy: Never
+  containers:
+  - name: bench-client
+    image: vllm/vllm-tpu:nightly-20260611-1043491-248e33c
+    command: ["/bin/bash", "-c"]
+    args:
+    - |
+      echo "Waiting for vLLM server to be ready..." && \
+      while ! curl http://vllm-gemma4-service:8000/ping; do sleep 5; done && \
+      echo "Downloading ShareGPT dataset..." && \
+      curl -L -o /tmp/sharegpt.json https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json && \
+      echo "Server is ready! Launching ShareGPT serving benchmark..." && \
+      vllm bench serve \
+        --backend=vllm \
+        --host=vllm-gemma4-service \
+        --port=8000 \
+        --model=google/gemma-4-31B-it \
+        --tokenizer=google/gemma-4-31B-it \
+        --dataset-name=sharegpt \
+        --dataset-path=/tmp/sharegpt.json \
+        --num-prompts=100 \
+        --request-rate=10 \
+        --percentile-metrics=ttft,tpot,itl,e2el \
+        --temperature=0.0 \
+        --ignore-eos
+    env:
+    - name: HF_TOKEN
+      valueFrom:
+        secretKeyRef:
+          name: hf-secret
+          key: hf_api_token
diff --git a/inference/trillium/vLLM/Gemma4-MTP/gemma4-mtp-gke.yaml b/inference/trillium/vLLM/Gemma4-MTP/gemma4-mtp-gke.yaml
@@ -48,6 +48,8 @@ metadata:
   namespace: gemma4-mtp
 spec:
   replicas: 1
+  strategy:
+    type: Recreate
   selector:
     matchLabels:
       app: vllm-gemma4-server
@@ -57,7 +59,7 @@ spec:
         app: vllm-gemma4-server
     spec:
       nodeSelector:
-        cloud.google.com/gke-tpu-accelerator: tpu-v6e
+        cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
         cloud.google.com/gke-tpu-topology: 2x2
       containers:
       - name: vllm-server
@@ -118,7 +120,7 @@ spec:
           value: "1"
         - name: XLA_PYTHON_CLIENT_MEM_FRACTION
           value: "0.75"
-        - name: XLA_CACHE_DIR
+        - name: VLLM_XLA_CACHE_PATH
           value: "/data/xla_cache"
         - name: VLLM_TORCH_COMPILE_CACHE_DIR
           value: "/data/torch_compile_cache"
diff --git a/inference/trillium/vLLM/Gemma4-MTP/recipe/Chart.yaml b/inference/trillium/vLLM/Gemma4-MTP/recipe/Chart.yaml
@@ -0,0 +1,6 @@
+apiVersion: v2
+name: v6e-benchmark
+description: v6e-benchmark
+type: application
+version: 0.1.0
+appVersion: "1.16.0"
diff --git a/inference/trillium/vLLM/Gemma4-MTP/recipe/README.md b/inference/trillium/vLLM/Gemma4-MTP/recipe/README.md
@@ -0,0 +1,92 @@
+# Run vLLM gemma-4-31B-it on v6e on GKE
+
+This recipe covers running a vLLM inference workload on v6e on
+GKE.
+
+## Create the GKE Cluster
+Create your v6e cluster using [XPK](https://github.com/AI-Hypercomputer/xpk).
+The next sections assume you have created a cluster with v6e
+nodes.
+
+## Deploy vLLM Workload on GKE
+
+### Configure kubectl to communicate with your cluster
+
+```
+gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${LOCATION}
+```
+
+### Generate a new Hugging Face token if you don't already have one
+
+On the huggingface website, create an account if necessary, and go to Your
+Profile > Settings > Access Tokens.
+Select Create new token.
+Specify a name of your choice and a role with at least Read permissions.
+Select Generate a token, follow the prompts and save your token.
+
+(NOTE: Also ensure that your account has access to the model on Hugging Face.
+For example, for llama3, you will need to explicitly get permission):
+
+### Run the benchmark
+In this directory, run:
+
+```
+helm install ${RUN_NAME} . --set hf_token=${HF_TOKEN}
+```
+
+The benchmark will launch a client and server pod.
+
+On the server pod, at the end of the server startup you'll see logs such as:
+
+```
+$ kubectl logs deployment/vllm-tpu -f
+
+(APIServer pid=1) INFO:     Started server process [1]
+(APIServer pid=1) INFO:     Waiting for application startup.
+(APIServer pid=1) INFO:     Application startup complete.
+```
+
+The client pod will wait until the server is up and then start the benchmark.
+On the client pod, you'll see logs such as:
+
+```
+$ kubectl logs -f vllm-bench
+
+============ Serving Benchmark Result ============
+Successful requests:                     10
+Failed requests:                         0
+Benchmark duration (s):                  xx
+Total input tokens:                      xxx
+Total generated tokens:                  xxx
+Request throughput (req/s):              xx
+Output token throughput (tok/s):         xxx
+Peak output token throughput (tok/s):    xxx
+Peak concurrent requests:                10.00
+Total Token throughput (tok/s):          xxx
+---------------Time to First Token----------------
+Mean TTFT (ms):                          xxx
+Median TTFT (ms):                        xxx
+P99 TTFT (ms):                           xxx
+-----Time per Output Token (excl. 1st token)------
+Mean TPOT (ms):                          xxx
+Median TPOT (ms):                       xxx
+P99 TPOT (ms):                           xxx
+---------------Inter-token Latency----------------
+Mean ITL (ms):                           xxx
+Median ITL (ms):                         xxx
+P99 ITL (ms):                            xxx
+==================================================
+```
+
+### Customizing the benchmark
+
+In order to change the parameters of the benchmark, such as the model, number of
+prompts, etc, you can modify the settings in values.yaml.
+
+### Cleanup
+When you are done running the benchmark, the server will still be running. You
+can cleanup by running:
+
+```
+helm uninstall ${RUN_NAME}
+```
diff --git a/inference/trillium/vLLM/Gemma4-MTP/recipe/templates/benchmark.yaml b/inference/trillium/vLLM/Gemma4-MTP/recipe/templates/benchmark.yaml
diff --git a/inference/trillium/vLLM/Gemma4-MTP/recipe/values.yaml b/inference/trillium/vLLM/Gemma4-MTP/recipe/values.yaml
