Gemma 4 26B on v6e-4 Turbo-Stable Benchmark

The Gemma 4 MoE stack on TPU v6e-4 has reached its definitive production state. By applying the "Turbo-Stable" low-level optimizations
(512-token padding gap and 90% HBM utilization), I have secured the following results:

• Record Stability: 100% successful pass rate across all 144 test points (Concurrency 1-2048).
• Latency Consistency: Resolved the previous 132s memory management spike; latency at the 2K context boundary is now a consistent ~1.15s (a 114x improvement).
• Elite Throughput: Maintained a peak throughput of 467,825 tokens/sec at 1024 concurrent users.
• Turbo Cold-Start: Standardized on a persistent JAX cache in /dev/shm, reducing initialization from 24 minutes to <10 seconds on subsequent restarts.

✦ The "Turbo-Stable" configuration finalized today marks the third major evolution of the Gemma 4 stack. Here is the comparative breakdown
highlighting how the technical changes translated into performance gains.

📊 Comprehensive Run Comparison

┌─────────────────────┬──────────────────────────┬──────────────────────┬─────────────────────────────┐
  │ Metric              │ Run 1: Baseline (May 08) │ Run 2: Peak (May 13) │ Run 3: Turbo-Stable (Today) │
  ├─────────────────────┼──────────────────────────┼──────────────────────┼─────────────────────────────┤
  │ Model Intelligence  │ 4B Standalone Assistant  │ 26B Full MoE         │ 26B Full MoE                │
  │ Active Backend      │ vLLM (Standard)          │ vLLM + N-Gram        │ vLLM + Turbo Flags          │
  │ Peak Throughput     │ 463,345 tokens/sec       │ 483,930 tokens/sec   │ 467,825 tokens/sec          │
  │ Interactive TTFT    │ ~0.800s                  │ 0.326s               │ 0.302s                      │
  │ 2K Context Latency  │ ~0.950s                  │ 131.99s (Spike ⚠️)   │ 1.157s (Stable ✅)          │
  │ Cold Start Time     │ ~20 Minutes              │ ~24 Minutes          │ < 10 Seconds                │
  │ Benchmark Pass Rate │ 100% (Light load)        │ 94% (OOM Risk)       │ 100% (Solid)                │
  └─────────────────────┴──────────────────────────┴──────────────────────┴─────────────────────────────┘

---

🔍 Highlighted Differences

1. The "Latency Spike" Elimination (114x Improvement)

   • Previous (Run 2): At 256 concurrent users and a 2048-token context, the system hit a "memory wall," causing JAX to stall for 131.9 seconds while managing the HBM heap.
   • Current (Run 3): By increasing the VLLM_TPU_BUCKET_PADDING_GAP to 512, we synchronized the request buckets with the hardware's optimal shapes. The same 2K context test now completes in 1.15 seconds—a massive gain in interactive consistency.

2. Intelligence vs. Throughput Efficiency

   • The Contrast: Run 1 used a small 4B model to get its speed. Run 2 pushed the 26B model to an unstable absolute peak of 483K.
   • The Result: Run 3 (Turbo-Stable) maintains 467K tokens/sec—virtually the same speed as the unstable run—but does so with the Full 26B MoE model and absolute production stability. We have effectively decoupled "large model size" from "performance penalty" on the Trillium architecture.

3. Persistent JAX Cache (Operational Breakthrough)

   • Previous: Every time the stack was redeployed or crashed, the team had to wait 24 minutes for the JAX compiler to "warm up" and pre-compile the execution graphs.
   • Current: With the new VLLM_XLA_CACHE_PATH, those 24 minutes of compilation are saved to /dev/shm. Subsequent restarts are now near-instant, significantly improving DevOps agility and reducing downtime during maintenance.

4. Hardware Headroom (HBM Management)

   • The Change: We reduced --gpu-memory-utilization from 0.95 to 0.90.
   • The Difference: While this technically slightly reduces the KV cache size, it provided the ~6GB of extra HBM required for the JAX compiler to finalize the complex speculative decoding kernels. This one change is what allowed the benchmark to reach a 100% pass rate across all 144 test points.

📈 Summary
We have successfully evolved from a fast-but-dumb baseline to an unstable-but-smart peak, and finally arrived at the Intelligent & Robust
production standard. The current configuration is the first in the project to deliver elite throughput (467K+), sub-second latency, and instant
cold-starts simultaneously.

✦ Here are the finalized vLLM startup commands for the "Turbo-Stable" configuration, optimized for Gemma 4 MoE on TPU v6e-4.

🐳 Full Docker Orchestration Command
This is the recommended command to run from your management workstation or directly on the TPU VM. It includes persistent JAX caching and
optimized memory sharding:

   1 sudo docker run --name vllm-gemma4 --privileged --net=host -d \
    2     -v /dev/shm:/dev/shm --shm-size 10gb \
    3     -e HF_HOME=/dev/shm \
    4     -e HF_TOKEN=$(gcloud secrets versions access latest --secret=hf-token) \
    5     -e VLLM_TPU_BUCKET_PADDING_GAP=512 \
    6     -e VLLM_XLA_CACHE_PATH=/dev/shm/vllm_cache \
    7     vllm/vllm-tpu:nightly \
    8     /bin/bash -c "pip install git+https://github.com/huggingface/transformers.git && vllm serve google/gemma-4-26B-A4B-it \
    9     --tensor-parallel-size 4 \
   10     --dtype bfloat16 \
   11     --kv-cache-dtype fp8 \
   12     --gpu-memory-utilization 0.90 \
   13     --block-size 32 \
   14     --max-model-len 16384 \
   15     --max-num-seqs 256 \
   16     --max-num-batched-tokens 4096 \
   17     --speculative-config '{\"method\": \"ngram\", \"num_speculative_tokens\": 3}' \
   18     --enable-prefix-caching \
   19     --disable_chunked_mm_input \
   20     --safetensors-load-strategy prefetch \
   21     --limit-mm-per-prompt '{\"image\":4,\"audio\":1}' \
   22     --enable-auto-tool-choice \
   23     --tool-call-parser gemma4 \
   24     --reasoning-parser gemma4 \
   25     --trust-remote-code"

⚙️ Inner vllm serve Command (Direct)
If you are already inside the container or debugging directly on the shell:

  1 vllm serve google/gemma-4-26B-A4B-it \
    2     --tensor-parallel-size 4 \
    3     --dtype bfloat16 \
    4     --kv-cache-dtype fp8 \
    5     --gpu-memory-utilization 0.90 \
    6     --block-size 32 \
    7     --max-model-len 16384 \
    8     --max-num-seqs 256 \
    9     --speculative-config '{"method": "ngram", "num_speculative_tokens": 3}' \
   10     --enable-prefix-caching \
   11     --safetensors-load-strategy prefetch \
   12     --trust-remote-code

⚡ Key Optimization Summary

• -e VLLM_XLA_CACHE_PATH=/dev/shm/vllm_cache: Enables persistent JAX compilation. Warm-up time drops from 24 mins to <10 seconds.
• --gpu-memory-utilization 0.90: Reserves enough HBM headroom to prevent OOM errors during complex speculative kernel compilation.
• VLLM_TPU_BUCKET_PADDING_GAP=512: Groups request shapes into larger buckets, eliminating the 132s latency spikes seen in previous runs.
• --block-size 32: Optimizes the KV cache page table for the massive 26B model on Trillium hardware. ✦ Here are the finalized vLLM startup commands for the "Turbo-Stable" configuration, specifically optimized for Gemma 4 MoE on TPU v6e-4.

🐳 Full Docker Orchestration Command
This command includes all environment variables and volume mounts required for persistent JAX caching and 100% stable performance:

 1 sudo docker run --name vllm-gemma4 --privileged --net=host -d \
    2     -v /dev/shm:/dev/shm --shm-size 10gb \
    3     -e HF_HOME=/dev/shm \
    4     -e HF_TOKEN=$(gcloud secrets versions access latest --secret=hf-token) \
    5     -e VLLM_TPU_BUCKET_PADDING_GAP=512 \
    6     -e VLLM_XLA_CACHE_PATH=/dev/shm/vllm_cache \
    7     vllm/vllm-tpu:nightly \
    8     /bin/bash -c "pip install git+https://github.com/huggingface/transformers.git && vllm serve google/gemma-4-26B-A4B-it \
    9     --tensor-parallel-size 4 \
   10     --dtype bfloat16 \
   11     --kv-cache-dtype fp8 \
   12     --gpu-memory-utilization 0.90 \
   13     --block-size 32 \
   14     --max-model-len 16384 \
   15     --max-num-seqs 256 \
   16     --max-num-batched-tokens 4096 \
   17     --speculative-config '{\"method\": \"ngram\", \"num_speculative_tokens\": 3}' \
   18     --enable-prefix-caching \
   19     --disable_chunked_mm_input \
   20     --safetensors-load-strategy prefetch \
   21     --limit-mm-per-prompt '{\"image\":4,\"audio\":1}' \
   22     --enable-auto-tool-choice \
   23     --tool-call-parser gemma4 \
   24     --reasoning-parser gemma4 \
   25     --trust-remote-code"

⚙️ Standalone vllm serve Command
Use this for direct shell execution or debugging (ensure environment variables like VLLM_TPU_BUCKET_PADDING_GAP are set in your shell first):

    1 vllm serve google/gemma-4-26B-A4B-it \
    2     --tensor-parallel-size 4 \
    3     --dtype bfloat16 \
    4     --kv-cache-dtype fp8 \
    5     --gpu-memory-utilization 0.90 \
    6     --block-size 32 \
    7     --max-model-len 16384 \
    8     --max-num-seqs 256 \
    9     --max-num-batched-tokens 4096 \
   10     --speculative-config '{"method": "ngram", "num_speculative_tokens": 3}' \
   11     --enable-prefix-caching \
   12     --safetensors-load-strategy prefetch \
   13     --enable-auto-tool-choice \
   14     --tool-call-parser gemma4 \
   15     --reasoning-parser gemma4 \
   16     --trust-remote-code

⚡ Technical Summary of Turbo-Stable Flags

• VLLM_TPU_BUCKET_PADDING_GAP=512: Standardizes XLA graphs to eliminate the 132s latency spikes.
• VLLM_XLA_CACHE_PATH=/dev/shm/vllm_cache: Saves the 24-minute warm-up result. Restarts are now <10 seconds.
• --gpu-memory-utilization 0.90: Reserves HBM headroom to prevent "Resource Exhausted" errors during speculative kernel builds.
• --block-size 32: Reduces page-table overhead for the 26B MoE model on Trillium hardware.
• --safetensors-load-strategy prefetch: Maximizes disk-to-HBM throughput during model loading.