Repository navigation
Conversation
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Adds configurable diffusion continuous-batching support (compute-budget + profiling) and improves mixed-shape scheduling fairness by grouping waiting requests by sampling key.
Changes:
- Add CLI flags for diffusion batching config/policy and plumb them into engine/stage args.
- Introduce diffusion batching policy plugins (fixed/compute_budget/profile) plus profiler tooling + docs.
- Rework diffusion scheduler waiting-queue handling using keyed queues + cost-aware DRR selection; expand tests accordingly.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| vllm_omni/entrypoints/cli/serve.py | Adds CLI flags to configure diffusion batching policy and JSON config. |
| vllm_omni/engine/async_omni_engine.py | Passes diffusion batching config/policy into default diffusion stage args. |
| vllm_omni/diffusion/worker/diffusion_model_runner.py | Emits profiler log lines for denoise-step timing under profiler policy. |
| vllm_omni/diffusion/sched/batching_policy.py | New batching policy abstractions + implementations (fixed/compute budget/profiler). |
| vllm_omni/diffusion/sched/base_scheduler.py | Implements keyed waiting queues and DRR-based key selection; integrates batching policy cap. |
| vllm_omni/diffusion/data.py | Adds DiffusionBatchingConfig and wires it into OmniDiffusionConfig. |
| tests/diffusion/test_diffusion_scheduler.py | Updates/extends scheduler tests for keyed DRR + compute-budget/profile policies. |
| docs/user_guide/diffusion/compute_budget_batching.md | New user guide for compute-budget batching, profiling workflow, and mixed-shape scheduling. |
| benchmarks/diffusion/diffusion_batching_profiler.py | New script to derive compute-unit budgets from profiler logs. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| devices = ",".join(str(i) for i in range(num_devices)) | ||
| model_class_name = kwargs.get("model_class_name", None) | ||
| final_output_type = "audio" if model_class_name and supports_audio_output(model_class_name) else "image" | ||
| diffusion_batching_config = dict(kwargs.get("diffusion_batching_config") or {}) |
| key = self._select_next_waiting_key() | ||
| if key is None and key not in self._waiting_by_key: | ||
| return None | ||
| return self._peek_waiting_state_for_key(key) | ||
|
|
| # All non-empty queues were visited but did not have enough budget. | ||
| self._drr_cursor_index = start | ||
| return None |
| def _peek_waiting_state_for_key(self, key: SamplingParamsKey | None) -> DiffusionRequestState | None: | ||
| queue = self._waiting_by_key.get(key) | ||
| while queue: | ||
| state = self._request_states.get(queue[0]) | ||
| if state is not None and not state.is_finished(): | ||
| return state | ||
| queue.popleft() | ||
| if key in self._waiting_by_key: | ||
| self._waiting_by_key.pop(key, None) | ||
| return None |
| def _peek_waiting_state_for_key(self, key: SamplingParamsKey | None) -> DiffusionRequestState | None: | ||
| queue = self._waiting_by_key.get(key) | ||
| while queue: | ||
| state = self._request_states.get(queue[0]) | ||
| if state is not None and not state.is_finished(): | ||
| return state | ||
| queue.popleft() | ||
| if key in self._waiting_by_key: | ||
| self._waiting_by_key.pop(key, None) | ||
| return None | ||
|
|
||
| def _pop_waiting_state(self, state: DiffusionRequestState) -> None: | ||
| self._remove_waiting_id(state.sched_req_id) | ||
|
|
||
| def _remove_waiting_id(self, sched_req_id: str) -> None: | ||
| self._waiting = deque(req_id for req_id in self._waiting if req_id != sched_req_id) | ||
| state = self._request_states.get(sched_req_id) | ||
| if state is None: | ||
| return | ||
| key = state.sampling_params_key | ||
| queue = self._waiting_by_key.get(key) | ||
| if queue is None: | ||
| return | ||
| self._waiting_by_key[key] = deque(req_id for req_id in queue if req_id != sched_req_id) | ||
| if not self._waiting_by_key[key]: | ||
| self._waiting_by_key.pop(key, None) |
Signed-off-by: xuan58 <xuanai2-c@my.cityu.edu.hk>
2f0cf1b to
48115f3
Compare
|
fix conflicts |
|
Hello, thank you for your contribution. We are going to remove most fields from SamplingParamsKey. You may need to adapt the changes for #3846 |
The current |
hsliuustc0106
added
diffusion
ChatGPT Codex Connector
Omni ReviewBot triage noteAutomated triage of commit
These are automated triage suggestions only — the final decision belongs to the maintainers. |
Omni ReviewBot: three questions on the performance claim@xuan58 this PR reads as a performance or value claim:
Before the full evidence checklist, three short questions:
When you answer, the evidence that settles it is: base and head SHA, hardware, model, workload, warm-up and repeat count, mean or percentiles with their spread, and a correctness/quality-equivalence signal; an end-to-end claim also needs stage attribution. |
Omni ReviewBot: no human activity for 15 days@xuan58 this pull request has had no human commit, comment or review since 2026-08-31. Per repository policy it may be closed if it stays inactive. To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline. |
Omni ReviewBot: no human activity for 23 days@xuan58 this pull request has had no human commit, comment or review since 2026-08-31. Please consider marking this PR as draft until work can resume. The author or a maintainer decides whether to change the PR state. To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline. |
hsliuustc0106
added
the
stale
ChatGPT Codex Connector
Summary
This PR optimizes diffusion step-wise continuous batching for request-shape- and hardware-dependent workloads.
It adds:
SamplingParamsKeyqueues to reduce head-of-line blocking in mixed-resolution traffic.Closes #3831.
Motivation
In vLLM-Omni 0.20.0, I validated the release version performance on a single Ascend 910B NPU with the Qwen/Qwen-Image model.
PR #2729 introduced step-wise continuous batching for diffusion models, and the PR benchmark showed performance gains on a single A100 GPU with
max-num-seqs=8. However, when running the same configuration on a single Ascend 910B NPU, directly usingmax-num-seqs=8does not bring a performance improvement. In my validation, it slightly regressed end-to-end benchmark duration and average latency.This suggests that the optimal
max-num-seqsfor diffusion continuous batching is hardware- and request-shape-dependent. A fixed value such as8may work well on A100, but is not necessarily optimal on Ascend 910B.Problems
I ran the text-to-image benchmark with:
0.20.0vllm-omniQwen/Qwen-Imagevbench1001024x1024508Benchmark command:
Baseline:
max-num-seqs=1Server command:
Result:
Continuous batching:
max-num-seqs=8Server command:
Result:
Compared with
max-num-seqs=1,max-num-seqs=8regressed benchmark duration by about 2.42% and average latency by about 4.17%.This suggests that a single fixed
max-num-seqsis not robust across hardware and request shapes. A good batch size for A100 is not necessarily a good batch size for Ascend 910B, and a good batch size for 512x512 requests is not necessarily a good batch size for 1024x1024 requests.Profiling-Based Batch Size
This PR adds a small policy layer for diffusion continuous batching:
fixed: preserves the existing fixedmax_num_seqsbehavior.profile: ramps the effective batch cap and logs[DiffusionBatchProfiler]records.compute_budget: derives the effective homogeneous batch cap from a profiled compute-unit budget.The profiling workflow is: first profile a reference request shape, find the batch-size sweet spot from execution-time behavior, then use that profiled parameter for normal serving.
For diffusion, the reference shape is 512x512. The profiler records:
benchmarks/diffusion/diffusion_batching_profiler.pyparses these logs, groups 512x512 records by batch size, filters outliers with IQR, uses median denoise-step time for each batch size, and applies a two-segment linear fit to locate the 512x512 sweet-spot batch size.The generated policy parameter is:
During serving, the scheduler estimates each request type by visual compute units:
Then it derives:
max_num_seqsremains a hard upper bound.For the built-in Qwen-Image reference profile on one Ascend 910B2 NPU:
The built-in
10240profile should only be used when the model, backend, hardware, and reference profiling shape match this setup. Other environments should run the profiler first and pass the generatedcompute_unit_budget.Example profiling run:
Generate the serving config:
Example serving command on Ascend NPU:
drr_arrival_window_sizeandlog_statsare optional. The former controls the sliding window used to estimate request arrival probabilities for DRR; the latter only enables verbose batching-policy logs.Keyed DRR Scheduling
Step-wise continuous batching still requires requests in one running batch to share the same
SamplingParamsKey. With a single FIFO waiting queue, a high-cost or incompatible request at the head can block later lower-cost requests, even when those later requests could form an efficient homogeneous batch.This PR changes the waiting side of the scheduler to keep per-key queues:
The scheduler then selects the next homogeneous queue with a cost-aware DRR policy:
cost_iis estimated from the request shape as visual compute units.prob_iis estimated online from a sliding arrival window.prob_i / cost_i.After a key is selected, the existing homogeneous batching rule is preserved: only requests with the same
SamplingParamsKeycan join the running batch. Thecompute_budgetpolicy then caps the batch size for that selected request type.Performance
I evaluated a mixed-resolution Qwen-Image workload on one Ascend 910B NPU. Requests arrive at high concurrency with a 50/50 mix of 512x512 and 1024x1024 requests, each using 50 inference steps.
Benchmark command:
Fixed continuous batching,
max-num-seqs=8:Profiled compute-budget batching plus keyed DRR scheduling:
The optimized strategy improves benchmark duration from
830.84sto679.94s, which is about an 18.16% speedup.Documentation
This PR adds
docs/user_guide/diffusion/compute_budget_batching.md, which documents:compute_unit_budget.