Skip to content

Supports batching in vLLM offline inference script#3932

Open
learning-to-play wants to merge 1 commit into
mainfrom
learning-to-play-patch-4
Open

Supports batching in vLLM offline inference script#3932
learning-to-play wants to merge 1 commit into
mainfrom
learning-to-play-patch-4

Conversation

@learning-to-play
Copy link
Copy Markdown
Collaborator

@learning-to-play learning-to-play commented May 18, 2026

Description

Support batching in vLLM offline inference script

This change enables batching in the vllm_decode.py offline inference script by duplicating the input prompt based on the per_device_batch_size configuration parameter.

Details:

  • Replaces the hardcoded single-prompt list with a list multiplied by per_device_batch_size.
  • Explicitly casts per_device_batch_size to int to prevent TypeError when the parameter defaults to a float (e.g., 12.0 in base.yml).
  • Allows users to benchmark vLLM offline inference throughput with larger batch sizes by overriding per_device_batch_size on the command line.

Tests

I verified this change by running the vllm_decode script on a TPU v5e-8 VM with per_device_batch_size set to different values, confirming that it batches correctly and generates outputs.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@codecov
Copy link
Copy Markdown

codecov Bot commented May 18, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants