gemma-4-31B-it-optimized

<!-- TODO: add the upstream model's `license:` value to the frontmatter above before publishing. -->

Gemma-4-31B-it — Tuned vLLM Serving Recipe

Docker build, launch script, and benchmark-reproduction script for serving google/gemma-4-31B-it with vLLM.

This repository does not contain model weights. vllm serve pulls them directly from the upstream repo above (and caches them) when the server starts.

Model Overview

• Dense, instruction-tuned multimodal model, 30.7B parameters, all active
• 60 layers, hidden size 5,376
• 32 query heads / 16 KV heads (sliding-window layers), 32 query heads / 4 KV heads (global layers)
• Native context 262,144 tokens; benchmarked text-only with image and audio inputs disabled

Environment

Optimization Changes

• FP8 (e4m3) KV cache
• --gpu-memory-utilization 0.95
• --max-num-seqs 128, --max-num-batched-tokens 16384
• Enabled prefix caching and chunked prefill
• --limit-mm-per-prompt '{"image": 0, "audio": 0}' skips vision/audio encoder allocation, freeing VRAM for KV cache on this text-only workload
• VLLM_ATTENTION_BACKEND=FLASHINFER

Docker

Installation

Build the image:

docker build -t your-org/gemma-4-31b-it-vllm:latest .

Run it:

docker run --gpus all --ipc=host -p 8000:8000 \
  -e HF_TOKEN="hf_xxx" \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  your-org/gemma-4-31b-it-vllm:latest optimized

Google's Gemma releases are typically gated on the Hub. If google/gemma-4-31B-it requires accepting a license, do so on the model page and pass HF_TOKEN; otherwise the -e HF_TOKEN=... line can be dropped.

Pass baseline instead of optimized as the final argument to run the unoptimized configuration used for the comparison below.

Baseline

vllm serve google/gemma-4-31B-it\

--host 0.0.0.0 \

--port 8000 \

 --max-model-len 24576 \

 --gpu-memory-utilization 0.92 \

 --limit-mm-per-prompt '{"image": 0, "audio": 0}' \

  --trust-remote-code

Optimized

export VLLM_ATTENTION_BACKEND=FLASHINFER

vllm serve google/gemma-4-31B-it \

--host 0.0.0.0 --port 8000 \

--tensor-parallel-size 1 \

--max-model-len 24576 \

--gpu-memory-utilization 0.95 \

--kv-cache-dtype fp8_e4m3 \

--max-num-seqs 128 \

 --enable-prefix-caching \

--enable-chunked-prefill \

 --max-num-batched-tokens 16384 \

--limit-mm-per-prompt '{"image": 0, "audio": 0}' \

--trust-remote-code

Validation

Once the container is up, confirm the server is healthy and serving:

curl -sf http://localhost:8000/health && echo "server is healthy"

curl -s http://localhost:8000/v1/completions \
  -H "Content-Type: application/json" \
  -d '{"model": "google/gemma-4-31B-it", "prompt": "The capital of France is", "max_tokens": 16}'

Files in this repository

Shell Script

Launch Server

./launch_server.sh optimized   # or: ./launch_server.sh baseline

This is also the image's default command — docker run ... <image> optimized runs the same thing. Inside a running container, invoke it with docker exec -it <container> ./launch_server.sh <mode>.

Reproduce Benchmark

This is a separate container invocation from "Launch Server" above — it overrides the entrypoint so run_benchmark.sh (not launch_server.sh) drives the container. It starts the baseline server itself, benchmarks it, shuts it down, starts the optimized server, benchmarks that, shuts it down, and prints the comparison table — nothing to copy in by hand:

docker run --gpus all --ipc=host \
  -e HF_TOKEN="hf_xxx" \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  --entrypoint ./run_benchmark.sh \
  your-org/gemma-4-31b-it-vllm:latest both

That prints the full table in the terminal when it finishes. Pass baseline or optimized instead of both to run just one side. The raw per-run JSON (benchmark_baseline.json, benchmark_optimized.json) and each server's startup log (server_baseline.log, server_optimized.log) are left in the container's /workspace if you want them too:

docker cp <container>:/workspace/benchmark_optimized.json .

Methodology (vllm bench serve):

• Dataset: random (synthetic), range ratio 0.8
• Input length: 409–3,686 tokens (nominal 2,048, uniform)
• Output length: 51–460 tokens (nominal 256, uniform)
• 1,000 measured prompts, 10 warm-up prompts (discarded, not saved)
• Request rate: unbounded (inf) — max concurrency is the only throttle
• Max concurrency: 128 (target concurrency; override with -e MAX_CONCURRENCY=...)

Results

1x H100-80GB, vLLM v0.20.0. This is exactly the table run_benchmark.sh both prints — running it is the only step required to reproduce it: