Google pairs Gemma with vLLM for a 23x speedup

Google just laid out a blueprint for running Gemma models inside Dataflow streaming pipelines using the open-source inference engine vLLM. It gives developers a way to deploy LLM-powered data pipelines with minimal boilerplate code while drastically driving down compute costs.

Running large models in a streaming pipeline has traditionally been a headache. Standard GPU batching groups inputs together, but language models generate outputs token by token. If a short query gets batched with an essay prompt, the GPU sits stalled waiting for the long prompt to finish before returning any results.

Continuous batching meets Dataflow

vLLM solves this with continuous batching. Instead of waiting for a batch to clear, it recomputes and updates the batch on every generated token. Completed outputs return instantly, and new incoming prompts join the queue on the fly, keeping the GPU cores busy around the clock.

The tricky part has always been orchestration. Dataflow typically spins up one independent process per CPU core. If each process tries to load a full Gemma checkpoint into memory, the machine immediately crashes with out-of-memory errors. Managing a shared inference server across custom multiprocessing topologies usually requires specialized infrastructure code.

Dataflow's model manager eliminates that setup. It provisions a single vLLM instance in a dedicated inference process. Other worker processes simply send records asynchronously into vLLM's request queue. Using Dataflow's RunInference transform, an end-to-end pipeline that reads data, generates completions with Gemma 2B, and writes out the results takes just five lines of pipeline code.

A 23x efficiency jump

To measure the real-world difference, Google ran a benchmark on a single worker equipped with an Nvidia T4 GPU, processing 10,000 prompts from the P3 dataset against the google/gemma-2b model.

A naive baseline pipeline using default fixed-size batching took 59.137 vCPU hours to complete the job. With vLLM and continuous batching enabled, that same workload finished in just 2.481 vCPU hours. That is an efficiency boost of more than 23 times.

There is one caveat: neither pipeline was specifically tuned. A baseline pipeline tuned with optimized batch sizes would perform better. But the point is that vLLM hits high throughput out of the box without manual tweaking. Switching to another open model only requires changing a single string in the model handler. The full tutorial and reproducible notebook are live on Google Cloud's documentation site.

Why it matters

If your team runs batch or streaming ETL pipelines on Apache Beam and wants to embed open-source LLMs for categorization, extraction, or translation, this setup is a no-brainer to try. It abstracts away messy multiprocessing plumbing and lets engineers treat LLM inference like any other standard data transformation.

We think this is a smart move by Google to make its data infrastructure friendlier to the open-source AI stack. The main trade-off to keep in mind is that continuous batching optimizes heavily for aggregate throughput rather than ultra-low single-request latency, making it ideal for data processing pipelines rather than live interactive chat.