[rollout, vllm_omni] feat: add batched diffusion generate API (one B=N forward instead of N B=1)

[SamitHuang]

Summary

Adds vLLMOmniHttpServer.generate_batched so a single Ray request can ask the vLLM-Omni diffusion engine to produce num_outputs_per_prompt samples in one QwenImagePipelineWithLogProb forward pass — the same way diffusers.QwenImagePipeline handles num_images_per_prompt=N.

Existing generate behavior is preserved bit-for-bit: it now delegates to a shared _generate_engine_call helper with num_outputs_per_prompt=1 and returns the first DiffusionOutput.

Motivation

In the current FlowGRPO Qwen-Image rollout (e.g. examples/flowgrpo_trainer/run_qwen_image_ocr.sh with rollout.n=16):

• The agent loop fans out N concurrent generate calls per prompt.
• Each call submits a B=1 request to vllm-omni.
• AsyncOmniDiffusion then serializes them via ThreadPoolExecutor(max_workers=1), so the GPU effectively runs N sequential B=1 diffusion forwards.
• Plain diffusers runs a single B=N forward and ends up faster on the same workload despite using a slower attention backend.

generate_batched lets the rollout client issue a single request that asks the pipeline to do one B=N transformer forward, then receive N DiffusionOutput objects back.

Changes

• verl_omni/workers/rollout/vllm_rollout/vllm_omni_async_server.py

  • Factored shared engine-call logic into private _generate_engine_call (takes num_outputs_per_prompt).
  • generate is now a thin wrapper around it with num_outputs_per_prompt=1 → fully backwards compatible.
  • New generate_batched returns list[DiffusionOutput] by setting OmniDiffusionSamplingParams.num_outputs_per_prompt=N.
  • New _unpack_batched_result slices the batched OmniRequestOutput (PIL image list + batched custom-output tensors: latents / log-probs / prompt-embeds / masks) into per-sample DiffusionOutputs, preserving the contract of the existing single-image path.
  • Ruff isort re-grouped the vllm_omni.* / vllm.* imports below verl.* (auto-applied by the project's pre-commit config).

• scripts/bench_qwen_image_rollout.py (new)

  • Apples-to-apples T2I benchmark comparing vLLM-Omni and diffusers.QwenImagePipeline at 512×512.
  • Workload, dtype, max_sequence_length, num_inference_steps, and true_cfg_scale lifted from examples/flowgrpo_trainer/run_qwen_image_ocr.sh; resolution overridden to 512 per the original task.
  • --vllm-omni-mode {concurrent,batched} toggles between the legacy N × B=1 submission pattern and the new 1 × B=N path through generate_batched.
  • Detects and logs the active diffusion attention backend for each engine (vllm-omni: DIFFUSION_ATTENTION_BACKEND selection; diffusers: _AttentionBackendRegistry._active_backend + SDPA dispatch info).
  • Reports setup / warmup / steady-state throughput and peak GPU memory; supports --backend {vllm_omni,diffusers,both} and --output bench_results.json.

Compatibility

• No CLI / config / public-method signature changes for existing call sites. generate still returns DiffusionOutput and still runs as B=1.
• generate_batched is a new method; existing callers are not affected.
• No changes to vllm_omni, verl, trainer, or worker configs.

Test plan

• pre-commit run --files verl_omni/workers/rollout/vllm_rollout/vllm_omni_async_server.py scripts/bench_qwen_image_rollout.py — all hooks pass except check-naming-conventions, which fails on the unchanged tree (its grep excludes venv but not .venv site-packages installed at the repo root). Skipped for this commit only and unrelated to this PR.

• Benchmark sanity run on this workstation (Qwen-Image, 512×512, num_inference_steps=50, rollout_n=16, num_prompts=2 → 32 images, bf16, FA3 for vllm-omni / TORCH_SDPA for diffusers):

  Backend	Throughput	Note
  vllm-omni --vllm-omni-mode concurrent	0.403 img/s	16 × B=1, matches current agent-loop path
  vllm-omni --vllm-omni-mode batched	0.655 img/s	1 × B=16 via new generate_batched
  diffusers	0.528 img/s	1 × B=16

  vllm-omni batched is +62% vs vllm-omni concurrent and +24% vs diffusers on this workload.

• Wire generate_batched into the FlowGRPO agent loop / rollout client (follow-up PR; this change only exposes the API).

Duplicate-work check

Per AGENTS.md, searched open PRs on verl-project/verl-omni:

• gh api search/issues -f q="repo:verl-project/verl-omni is:pr is:open vllm-omni rollout batch" → 4 hits ([diffusion, rollout, trainer] feat: add BAGEL FlowGRPO support verl-project/verl-omni#66 BAGEL FlowGRPO, [Bugfix] Enable step-wise execution verl-project/verl-omni#81 step-wise execution, [Algo] DPO (online) training with SD3.5-medium verl-project/verl-omni#77 SD3.5 DPO, [trainer, diffusion] feat: add z-image support for flowgrpo training verl-project/verl-omni#76 z-image FlowGRPO). None of them touch the diffusion server's request-batching API; they are model/algo/feature additions on adjacent surfaces. No duplicate work.
• gh pr list --repo SamitHuang/verl-omni --state open → 0 open PRs.

AI assistance disclosure

This change was prepared with AI assistance. The submitter has reviewed every modified line, defends the design end-to-end, and ran the benchmark above to confirm the speedup.

[gemini-code-assist]

Code Review

This pull request introduces a performance benchmarking script for Qwen-Image text-to-image rollout and implements batched generation support in the vLLMOmniHttpServer to allow multiple outputs per prompt in a single forward pass. Feedback identifies a potential division-by-zero error in the benchmarking script's summarization logic when no iterations are executed.

[gemini-code-assist]

The _summarize function is susceptible to a ZeroDivisionError and a statistics.StatisticsError if the times list is empty. This can occur if the user sets --iters 0 via the command line. Adding a guard clause to handle empty results would make the benchmark script more robust.

Suggested change

	mean_t = sum(times) / len(times)
	med_t = statistics.median(times)
	if not times:
	return 0.0, 0.0, 0.0, 0.0
	mean_t = sum(times) / len(times)
	med_t = statistics.median(times)