neural-nova/gemma-4-31B-it-optimized

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

Hardware	1x NVIDIA H100-80GB
Input Length	409–3,686 tokens, nominal 2,048
Output Length	51–460 tokens, nominal 256
Framework	vllm bench serve
Base Docker image	vllm/vllm-openai:v0.20.0
vLLM version	v0.20.0 · Transformers 5.x
Parallelism	Tensor-parallel 1
Features	tensor parallel 1 · FlashInfer · FP8 KV cache · prefix caching · chunked prefill
Prompts Measured	1,000 prompts, 10 warm-ups excluded
Request Rate	inf
Tested Concurrency	128

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

1. 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

File	Purpose
Dockerfile	vllm/vllm-openai:v0.20.0-based image with the two scripts below baked in
launch_server.sh	Starts vllm serve with the baseline or optimized flags
run_benchmark.sh	Starts baseline + optimized servers, benchmarks each, and prints the comparison table below — fully automated
README.md	This model card

2. 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:

Metric	Baseline	Optimized
Request thoughput (req/s)	0.85	1.41
Output tok/s	220.3	365.6
Total tok/s	2013.7	3341.8
TTFT mean (ms)	133009.7	70845.6
TTFT p99 (ms)	155345.7	81829.6
TPOT mean (ms)	34.7	58.2
ITL mean (ms)	34.6	58.1
E2E mean (ms)	142003.6	85936.3
E2E median (ms)	150279.6	89032.0
Duration (s)	1178.6	710.2