feat: Support dynamic batching for heterogeneous diffusion requests

[asukaqaq-s]

What does this PR do?

related:
#874
#3846

plans:

• Validate the single-GPU path.
• Add multi-GPU parallelism compatibility.
• Add compatibility with acceleration features such as CUDA Graph.

This PR adds dynamic batching support for heterogeneous diffusion execute_stepwise requests.

It allows diffusion requests with different image resolutions, image token lengths, text lengths, denoise progress, and request-local CFG scalar values to be batched in the same step, as long as they share the same execution mode.

Qwen-Image is used as the first supported model path to validate the metadata contract and execution flow.

Why is this needed?

diffusion batching is mostly limited by SamplingParamsKey. Several request-local fields are treated as hard batching keys, even though they do not change the transformer execution path.

For example, requests with different resolutions or different true_cfg_scale values may still be able to share one packed attention forward. However, because these fields are included in the compatibility key, the scheduler splits them into separate batches. This reduces batching opportunities and lowers GPU utilization.

This PR moves shape-related and CFG scalar fields out of the scheduler key and keeps them as request-local runner/model metadata.

Main changes

This PR:

• Narrows the dynamic diffusion batching key to true execution barriers:

  • do_classifier_free_guidance
  • lora_int_id
  • lora_scale

• Keeps resolution, text length, timestep, scheduler state, RoPE inputs, and CFG scalar values as request-local state.

• Adds attention metadata for:

  • request boundaries
  • text/image/negative-text spans
  • packed token lengths
  • request-local position/RoPE metadata

• Adds a FlashAttention varlen path for packed [total_tokens, num_heads, head_dim] Q/K/V tensors.

• Adds a Qwen-Image dynamic path that:

  • rebuilds per-request RoPE
  • runs packed joint attention
  • splits outputs back by request
  • applies CFG combine per request

• Keeps the existing dense homogeneous path unchanged.

• Fails fast for unsupported profiles.

How batching works

Instead of padding every request to the largest shape, the runner packs only real tokens.

For each step batch:

• image latents remain request-local before packing;
• text lengths are taken from the real prompt mask or request-local metadata;
• image grid shapes stay in img_shapes;
• packed spans are recorded in request/segment metadata;
• padded text or image tokens are not allowed to enter the dynamic attention path.

The packed path only changes how tokens are laid out for FlashAttention. It does not change request-local RoPE, attention boundaries, or output restore semantics.

Correctness model

The dynamic path relies on three metadata layers.

1. Request-local state

Each request keeps its own shape, text length, timestep, scheduler state, CFG values, and RoPE inputs.

RoPE is rebuilt from this request-local state. The model does not infer positions from packed global offsets, so a request should get the same RoPE whether it runs alone or inside a mixed batch.

2. AttentionMetadata

AttentionMetadata describes the varlen execution contract consumed by the attention backend.

It records packed request order, q_cu_seqlens, kv_cu_seqlens, max lengths, padded token count, and position metadata.

The FlashAttention varlen path consumes these boundaries directly, so tokens from one request cannot attend to tokens from another request even when they are concatenated into one flat tensor.

3. AttentionSegment

AttentionSegment records semantic split/gather spans inside each request.

It is used to identify text, image, and negative-text spans, and to scatter attention outputs back to the original request-local latent state.

AttentionSegment is intentionally separate from cache layout, LoRA routing, and MoE routing.

CFG behavior

CFG batching is split into two parts:

• execution shape
• request-local scalar math

The scheduler only keeps one CFG execution key: do_classifier_free_guidance.

In the first validated Qwen-Image path, a request is treated as CFG-active when:

true_cfg_scale > 1 and negative_prompt is present

This prevents CFG-active and non-CFG requests from sharing an incompatible forward shape.

However, different CFG scalar values can still batch together. For example, requests with true_cfg_scale=2.0 and true_cfg_scale=7.0 can share the same CFG batch because the scalar does not change the attention layout.

At runner/model time:

• the positive pass uses positive attention metadata;
• the negative pass uses negative attention metadata when true CFG is active;
• positive and negative text spans are represented as separate segments;
• true_cfg_scale and cfg_normalize remain per-request values;
• CFG combine is applied after per-request split:

negative + true_cfg_scale * (positive - negative)

This ensures that different guidance scales can batch together without mixing request semantics.

Current guardrails

This PR keeps the first dynamic path conservative.

The dynamic path currently fails fast for:

• KV cache / cache-enabled paths
• image editing inputs
• sequence parallel
• ring attention
• CFG parallel
• HSDP
• non-eager graph paths
• non-FlashAttention backends

These are intentionally disabled until their metadata and correctness requirements are handled explicitly.

Test plan

Focused scheduler and Qwen-Image dynamic batching tests:

python -m pytest \
  tests/diffusion/test_qwen_image_dynamic_step_batching.py \
  tests/diffusion/test_diffusion_scheduler.py

Online serving e2e test:

python -m pytest \
  tests/e2e/online_serving/test_qwen_image_dynamic_step_batching.py

Lint and format checks for changed Python files:

CHANGED_PY_FILES="$(git diff --name-only origin/main...HEAD -- '*.py')"
.venv/bin/ruff format $CHANGED_PY_FILES
.venv/bin/ruff check $CHANGED_PY_FILES

Test results

Focused scheduler and dynamic batching tests:

57 passed

Online serving e2e test:

1 passed

The current e2e coverage uses Qwen-Image and compares execute-model, dynamic single-request, and staggered dynamic multi-request paths for:

• 512x512 + 512x512
• 512x512 + 1024x1024
• 512x512 + 256x256

It also covers batching requests with different true_cfg_scale values and verifies that the dynamic multi-request path does not materially regress against the dynamic single-request path in the tested cases.

CC List.

cc
@ZJY0516 @fake0fan @SamitHuang @hsliuustc0106 @fhfuih @wuhang2014 @Semmer2 @wtomin

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[asukaqaq-s]

This PR has only been tested on single GPU so far. Multi-GPU and CUDA Graph validation still need to be added.
So this PR remains a draft for now.

[asukaqaq-s]

I also ran an online serving evaluation for the dynamic stepwise batching path. on one a100 card.
The evaluation compares execute-model serial outputs with execute-stepwise dynamic batching outputs. It covers three mixed-resolution cases, with six generated images in total.

Case	execute-model (serial)	dynamic batch	Speedup	max_abs	mean_abs	p99_abs
512+512	2.729807s	2.297938s	1.1879x	0	0.0	0.0
512+256	2.378886s	1.936803s	1.2283x	0	0.0	0.0
512+1024	5.290979s	5.222099s	1.0132x	0	0.0	0.0

The generated images from the dynamic batch path are pixel-wise identical to the execute-model serial reference in all tested cases.

Overall, this confirms that the current PR preserves output correctness for these mixed-resolution online serving cases. Performance does not regress. Same-size and small mixed-size batches show clear speedups, while the 512+1024 case is mostly dominated by the large request and therefore stays close to parity.

related pic

The two images above are generated in parallel: the first is 512×512 resolution, and the second is 1024×1024.

These two images are generated sequentially via the execute-model path.

[hsliuustc0106]

is this PR ready for review?

[asukaqaq-s]

is this PR ready for review?

It may take until next week to resolve the issue.

[desenSunUBW]

Hi, I compared this PR to #3638 on an Ada 6000 Pro GPU for Qwen-image model, and here is the latency results:

backend	resolution	mixfusion_serial_s	mixfusion_batched_s	pr4024_serial_s	pr4024_batched_s
Flash Attention	1024x1024,768x1024	4.37877	4.159676	4.231262	4.168515
Flash Attention	512x512,512x768	2.103336	1.798746	1.969378	1.927315

Maybe you can evaluate #3638 to A100 as well, to see if the conclusion remains.

[asukaqaq-s]

Summary

We added packed varlen step batching to Qwen-Image in vLLM-Omni. Previously, step batching only supported requests with the same resolution; now each denoise step can pack requests with different resolutions into a single [total_image_tokens, C] forward, using joint text+image varlen attention and packed GEMMs.

The batched output is bit-exact to the serial output, with max_abs = 0. Batching delivers around 1.1–1.8× speedup depending on the resolution mix and parallel mode, in both BF16 and INT8.

Setup

• Model: Qwen/Qwen-Image
• Hardware: 2× NVIDIA A100 80GB PCIe
• Stack: vLLM-Omni, based on vLLM 0.23.0, eager mode
• Steps: 4 denoise steps
• Batch: 2 concurrent requests
• Precision: BF16 and INT8
  --diffusion-quantization-config '{"method":"int8"}'
• Parallel modes:
  • single: 1 GPU
  • TP2: tensor parallel
  • USP2: Ulysses sequence parallel
  • Ring2: ring sequence parallel
• Determinism: step batching runs with VLLM_BATCH_INVARIANT=1, so packed-M GEMMs are numerically identical to per-request M GEMMs. This is what makes batch and serial outputs bit-exact.

What we measure

• serial: the two requests are denoised sequentially, step by step.
• batch: the two requests are packed into the same denoise step.
• Speedup: serial_wall / batch_wall, under the same build, precision, and parallel mode.
• Accuracy: element-wise max_abs diff between batched and serial outputs, measured at three points:
  • DiT noise: dit
  • final latent: latent
  • VAE image tensor: image_tensor

The threshold is ≤ 1e-5.

Accuracy

Across all configurations:

• BF16 / INT8
• single / TP2 / USP2 / Ring2
• 512+512 / 512+768 / 512+1024

The result is:

dit, latent, and image_tensor batch-vs-serial max_abs = 0, PASS for every cell.

Batching introduces zero numerical drift. The batched result is identical to the serial result.

Performance

BF16 speedup

mode	512+512	512+768	512+1024
single	1.25×	1.15×	1.12×
TP2	1.72×	1.35×	1.23×
USP2	1.63×	1.47×	1.30×
Ring2	1.54×	1.30×	1.22×

INT8 speedup

mode	512+512	512+768	512+1024
single	1.27×	1.15×	1.13×
TP2	1.74×	1.61×	1.37×
USP2	1.73×	1.58×	1.39×
Ring2	1.79×	1.58×	1.38×

Representative absolute wall time:

• BF16 TP2 512+512: serial 3.64s → batch 2.12s

Notes

1. in some bit-exactness case depends on deterministic batch-invariant GEMMs. With VLLM_BATCH_INVARIANT=1, batch and serial outputs reach max_abs = 0. Without it, batch and serial outputs typically match only to around 1e-3, due to normal kernel and reduction-order differences.

2. The speedups above are measured in eager mode with 4 denoise steps and 2 concurrent requests. Larger step counts and more concurrent requests generally improve the batching gain further.