[Perf][Diffusion][Qwen-Image] Replace per-bucket varlen loop with flat varlen attention for mixfusion

[renyongxing]

Purpose

Continue the work of #3638 (Efficient Mixed-Resolution Batching Strategy for HunyuanImage-3.0 and Qwen-Image). While fully preserving the original author's changes (standalone commit 43ffa56a9, authored by desenSunUBW), replace Qwen-Image mixfusion's per-bucket equal-length varlen loop with a single flat varlen attention pass to eliminate inter-bucket bubbles and reduce end-to-end latency for mixed-resolution batches.

Scope

Following review feedback that this PR was over the size budget, the non-Qwen-Image work has been split out:

• [Perf][Diffusion][Hunyuan] Bucket-based MixFusion for Hunyuan-Image-3.0 #8116 — [Perf][Diffusion][Hunyuan] Bucket-based MixFusion for Hunyuan-Image-3.0
• [Benchmark][Diffusion] MixFusion benchmarks #8117 — [Benchmark][Diffusion] MixFusion benchmarks

What remains here is the Qwen-Image flat-varlen change plus the shared scheduling/batching infrastructure it requires — 667 non-test added lines across 11 files. No Hunyuan-specific code and no benchmark tooling is left in this PR. The one shared piece both family PRs build on, _apply_mixfusion_shape_relaxation in vllm_omni/diffusion/sched/base_scheduler.py, intentionally stays here rather than being duplicated.

Key changes:

• vllm_omni/diffusion/attention/backends/abstract.py + flash_attn.py: AttentionMetadata gains varlen fields (is_varlen, q_cu_seqlens, kv_cu_seqlens, max_q_len, max_kv_len, padded_tokens); add _forward_varlen_flat (expects [total_tokens, H, D], calls flash_attn_varlen_func with the same interface as Enable Efficient Mixed-Resolution Batching Strategy for HunyuanImage-3.0 and Qwen-Image #3638).
• vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py: forward_mixfusion switches from the per-bucket loop to a single flat varlen pass (cu_seqlens built from real text lengths via mask.sum, text packed with 17-way padding, output unpacked by cu_seqlens).
• Fixes surfaced by the single-request API migration (see commits 4cde003a0, 28d3c272c): mixfusion benchmarks build one request per prompt + asyncio.gather; state_cache pop key uses request_id; OmniDiffusionRequest.prompts compatibility accessor for scheduler/pipeline req.prompts iteration.

Reference / Acknowledgement

The flat-varlen FlashAttention method used in this PR (varlen AttentionMetadata fields q_cu_seqlens / kv_cu_seqlens / max_q_len / max_kv_len /padded_tokens; flat packed [total_tokens, heads, head_dim] Q/K/V tensors; and the flash_attn_varlen_func-based branch that rejects attention masks and requires padded_tokens=0) is adapted from the approach proposed in PR #4024 — [Diffusion] Enable Qwen-Image dynamic step batching by @asukaqaq-s .

How the application differs:

• feat: Support dynamic batching for heterogeneous diffusion requests #4024 applies flat varlen attention to dynamic step batching — packing requests that are at different denoising steps into a single forward.
• This PR applies the same flat-varlen technique to mixed-resolution mixfusion batching (Enable Efficient Mixed-Resolution Batching Strategy for HunyuanImage-3.0 and Qwen-Image #3638), replacing the per-bucket equal-length varlen
  loop with one flat varlen pass over the mixfusion batch.

Note: PR #4024 is not yet merged into main. If it lands first, the overlapping AttentionMetadata varlen fields and the flat varlen branch in flash_attn.py should be reconciled to avoid duplication.

Test Plan

Automated tests

• Unit (packing/varlen layout): tests/diffusion/models/qwen_image/test_qwen_image_mixfusion_packing.py
  runs under pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' (the Simple · Diffusion & Model Executor Test selector).
• e2e: tests/e2e/online_serving/test_qwen_image.py -m core_model --run-level core_model
  (the Diffusion · Qwen Image Test selector) guards the Qwen-Image path this PR touches.
• Coverage gap (explicit): no CI e2e case currently passes enable_mixfusion=true; mixfusion itself
  is exercised by the packing unit tests above plus the manual benchmark runs below. The benchmark
  harness (mixfusion_benefit.py) now lives in [Benchmark][Diffusion] MixFusion benchmarks #8117.

1. Correctness (bitwise): run old (per-bucket) and new (flat varlen) implementations on the same 4 concurrent requests and assert latents are bitwise identical ((a == b).all() == True).
2. Benchmark matrix (A100, steps=8, iters=3, tp=1/2/4). The harness lives in [Benchmark][Diffusion] MixFusion benchmarks #8117; the command is unchanged apart from --family qwen:

   DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN \
   python benchmarks/diffusion/mixfusion_benefit.py \
     --family qwen \
     --model /data/weight/Qwen-Image \
     --image-sizes 1024x1024,1024x768,768x1024,512x512 \   # caseB (mixed resolution)
     --steps 8 --iters 3 --warmup 1 --max-num-seqs 4 \
     --tensor-parallel-size 1 --no-skip-rejected-mixfusion-cases

3. Smoke test (caseB, tp1): same benchmark with steps=8, iters=1.

vLLM Version: 0.28.0

**vLLM-Omni Commit:**0.28.0

Test Result

Scenario	batched time(Compared to the bucket version)	speedup	bubble
caseA (same resolution, tp1/2/4)	−2.3% ~ −7.3%	+0.05 ~ +0.12	—
caseB (mixed resolution, tp1/2/4)	−4.6% ~ −5.6%	+0.06 ~ +0.15	halved

• Correctness: 4 requests' latents bitwise identical between old and new implementations.
• Smoke test (caseB: 1024x1024/1024x768/768x1024/512x512, tp1, steps=8, iters=1, vllm 0.28.0 + vllm-omni 0.28.0 aligned env):
  • mixfusion plan accepted: chunk_size=1024, chunk_count=11
  • serial 10.62s / batched 9.30s → speedup 1.14
  • 4 images emitted in both serial and batched modes, backend FLASH_ATTN
• Original PR: Enable Efficient Mixed-Resolution Batching Strategy for HunyuanImage-3.0 and Qwen-Image #3638 (not yet merged upstream, kept as a standalone commit with authorship preserved)

The numbers above were produced with the pre-split per-family benchmark script; only the entry point moved, the measured code path is the one in this PR.

How to enable MixFusion

enable_mixfusion is opt-in per request, passed through the diffusion extra_args field:

{"extra_args": {"enable_mixfusion": true}}

It is accepted both at the top level of the request body and nested under
extra_body.extra_args. Two optional tuning knobs live in the same dict:

key	default	meaning
mixfusion_min_chunk_tokens	256	reject the flat-varlen layout when gcd(token counts) < this
mixfusion_max_chunks	128	reject when the resulting chunk count exceeds this

Online serving:

vllm serve Qwen/Qwen-Image --omni --port 8099

curl -s http://localhost:8099/v1/images/generations \
  -H 'Content-Type: application/json' \
  -d '{"model": "Qwen/Qwen-Image",
       "prompt": "a cat",
       "size": "1024x1024",
       "extra_args": {"enable_mixfusion": true}}'

Default is False, so nothing changes unless you ask for it

The key does not exist on main — it is introduced by this PR, and every read site defaults it to False.A deployment that never passes enable_mixfusion therefore behaves as before.

[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.

[renyongxing]

The MFU of the attention operator in this PR has not reached the ideal range, and the cause is still under investigation.

[vllm-omni-review-bot]

This PR appears to be related to model: qwen-image.

Model owners: @NumberWan

@renyongxing, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

[tzhouam]

@vllm-omni-review-bot please review

[vllm-omni-review-bot]

Omni ReviewBot: no human activity for 10 days

@renyongxing this pull request has had no human commit, comment or review since 2026-08-29. 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.

[renyongxing]

The Rope encoding in mixfusion is optimized by using the fusion operator. It is expected to be submitted and the comparison experiment result will be uploaded this week.

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit dea5597c0800 produced:

• Priority: high. Prompt maintainer attention is suggested.

These are automated triage suggestions only — the final decision belongs to the maintainers.

[renyongxing]

Hi @hsliuustc0106, thanks for the re-review and for confirming the earlier fixes. Both remaining items are addressed on the current head (0ccca59):

1. HSDP native-hook bypass — fixed in 0ccca59

has_active_block_hooks() now treats PyTorch-native hooks (block._forward_pre_hooks / block._forward_hooks) as active, in addition to Omni's HookRegistry. The pipeline re-checks this on every denoise step, and when any block carries an active hook it falls back to _denoise_step_serial_ragged, whose per-request loops route every block through the hooked module call — so fully_shard's all-gather pre-hooks run as usual.

We went with the serial fallback rather than register_fsdp_forward_method because the check is mechanism-agnostic: it protects against any native forward hook on the blocks, not just FSDP's, and needs no coupling to the HSDP implementation. The fallback also logs, so a silent packed-path loss is visible in server logs. Covered by new CPU unit tests that simulate fully_shard's pre-hook behavior (tests/diffusion/models/qwen_image/test_qwen_image_mixfusion_hooks.py).

2. Server-level enablement — ee69278

MixFusion previously required every request to opt in via sampling.extra_args["enable_mixfusion"]. There is now a server-level default in OmniDiffusionConfig (enable_mixfusion, mixfusion_min_chunk_tokens, mixfusion_max_chunks), propagated from the orchestrator CLI through StageDeployConfig. The per-request switch is still there and takes precedence: both the pipeline gate and the scheduler key read extra_args.get("enable_mixfusion", od_config.enable_mixfusion), so an explicit per-request True/False always overrides the server default.

[renyongxing]

@hsliuustc0106 To quantify the advantages of MixFusion in tensor parallelism and be closer to the actual scenario, we use main (step by step) in the case of TP △ {1, 2, 4}. MixFusion off) compared to b-on (Enable MixFusion at the server level). Experiment setup: Qwen-Image, 50 denoising steps, --enforce-eager, run 8 requests each time, and run each configuration for 3 times on average.

Two modes of arrival are reported:

• **s0 - Simultaneous Burst * *: All 8 requests are submitted simultaneously at t = 0, and the size is (1024²×2, 1024×768×2, 768×1024×2, 512²×2) in adjacent pairs of the same resolution. This is the most advantageous workload for main: requests of the same resolution share the same step-by-step batch key and can be scheduled simultaneously.
• s1 - Interleaved Heterogeneous Arrivals: One request is submitted every 2.5 seconds, (1024², 1024×768, 768×1024, 512², 1024², 1024×768, 768×1024, 512²) of sizes in mixed order, so requests of different sizes overlap on the fly. This represents real-world heterogeneous traffic.

Result - s0 (simultaneous burst, same size pairing):

TP	Main Total Duration (s)	MixFusion Total Duration (s)	Acceleration Ratio
1	90.96	83.33	1.09×
2	57.38	50.29	1.14 ×
4	40.83	32.42	1.26×

**Results - s1 (staggered arrivals, mixed size order): * *

TP	Main Total Duration (s)	MixFusion Total Duration (s)	Acceleration Ratio
1	95.42	84.73	1.13×
2	66.49	52.03	1.28 ×
4	50.23	35.38	1.42×

**Observations: * *

1. In all TPs, MixFusion is better than main, and the acceleration ratio increases with the increase of TPs (s0: 1.09× → 1.14× → 1.26×; s1: 1.13× → 1.28× → 1.42×). As TPs increase, each step of computation is spread over more GPUs, and execution is increasingly affected by scheduling/batch behavior, which magnifies the cost of `main' restricted batch processing.
2. **mainOnly requests with the same sampling parameters can be processed in batches. * * -- Step by step batch keys contain height and width (sched/interface.py). Its queue-head admission logic (base_scheduler.schedule()) blocks all subsequent requests after requests with mismatched keys. In s0 mode, it can still form pairs of the same size (; In s1 mode, the neighbors are all different sizes, so main degenerates to near-serial execution, with a total time that even exceeds the s0 baseline.
3. MixFusion is insensitive to the arrival sequence and size heterogeneity: It packs requests of different sizes into a single stepped batch and executes them step by step. Therefore, the total time consumed by the two arrival modes is basically the same. The largest difference is precisely where the flow heterogeneity is strongest (s1, TP=4: 1.42×).

**Summary: **The advantage of MixFusion comes from the packaging capability of heterogeneous requests, which cannot be achieved by the main key precise step-by-step batch processing. This advantage persists in tensor parallelism and is further magnified. In actual scenarios, the resolutions requested by users are not uniform, and inference is usually not performed on a single card. Therefore, Mixfusion can bring considerable benefits.

[hsliuustc0106]

Re-reviewed at 0ccca595ed977f7ea765e790531d1847e75af85d against base ee81f33001922f8d91dd09da0a57ee1531a8df82.

[P2] Register the server-level CLI flags — arg_utils.py:627

The new fields exist in OrchestratorArgs, but OmniServeCommand.subparser_init() never registers --enable-mixfusion, --mixfusion-min-chunk-tokens, or --mixfusion-max-chunks. Consequently, vllm serve … --omni --enable-mixfusion fails with an unrecognized-argument error before configuration propagation. Register the three flags, document their defaults, and add parser coverage.

The previous HSDP P1 is fixed: native PyTorch hooks now trigger serial fallback. No new correctness blocker found in the reviewed delta.

The remote head remained unchanged. All five reported checks and diff hygiene pass; no Buildkite result is listed. Validation was static—tests and GPU benchmarks were not independently rerun on the fork head.

[vllm-omni-review-bot]

Omni ReviewBot: superseded

The CI failure noted on f216a0589362 refers to an earlier head; the pull request now points at 86369e45af11.

[renyongxing]

@hsliuustc0106 Thanks for the reviews. A few updates since the last round:

Serve CLI flag registration (P2). Fixed in 5096155: --enable-mixfusion (BooleanOptionalAction, so --no-enable-mixfusion also works), --mixfusion-min-chunk-tokens (default 256), and --mixfusion-max-chunks (default 128) are now registered on the serve subparser and propagate through the stage-CLI override chain. Three parser tests cover registration and propagation.

Pipeline capability gating for MixFusion (new, f4d92ff). In self-review we found that the shape relaxation in _apply_mixfusion_shape_relaxation was gated only on the enable_mixfusion flag, while only QwenImagePipeline sets use_ragged_latents. On any other step-execution pipeline, server-level --enable-mixfusion or a request-level flag would relax height/width/resolution out of the batch key, co-batch mixed resolutions, and crash per step in the dense _prepare_latents path. This is now gated end-to-end on an opt-in pipeline capability:

• Pipelines declare supports_mixfusion: ClassVar[bool] = True (currently Qwen-Image only). The attribute is read via getattr and documented on SupportsStepExecution like supports_chunk_step_grouping, so the @runtime_checkable protocol keeps matching every pipeline.
• registry.supports_mixfusion(od_config) resolves the selected pipeline class (custom pipeline_class → Diffusers adapter → registry) without instantiating it, and the engine rejects server-level enable_mixfusion on unsupported pipelines at startup.
• The schedulers resolve the capability in initialize(). A request-level extra_args={"enable_mixfusion": True} on an unsupported pipeline degrades to a no-op — resolution stays in the batch key, so mixed resolutions keep batching separately. MixFusion is performance-only, so the degradation is safe.
• Tests cover capability resolution (declared / not declared / unresolvable), both schedulers' degradation paths, and the engine startup check; the affected suites pass 100 tests.

[vllm-omni-review-bot]

Omni ReviewBot: superseded

The CI failure noted on 86369e45af11 refers to an earlier head; the pull request now points at 08ce4fcb106a.

[vllm-omni-review-bot]

Omni ReviewBot: superseded

The CI failure noted on 051bff8517ef refers to an earlier head; the pull request now points at 241882deb03d.

[renyongxing]

@hsliuustc0106 I conducted a Buildkite verification locally, and the results are as follows.

Local reproduction of the Buildkite "Simple Test" (cuda/test-ready.yml) — failure analysis and re-trigger request

Setup

• PR branch head: dea5597c0 (tests executed at 241882deb; the only diff to the current head is the merged TTS test file from [Bugfix] Calibrate Qwen3-TTS PCM speech HNR floor for stochastic sampling #9029, so the conclusions carry over)
• Baseline: main @ f6272d885
• Same 4 groups as the CI config: G1 tests/model_executor, G2 tests/diffusion/models, G3 entrypoints + engine + diffusion, G4 the rest

Raw numbers

Group	PR run (FAILED/ERROR)	Baseline run (FAILED)
G1	38 / 7	24
G2	91 / 4	2
G3	162 / 12	8
G4	52 / 45	0

The two runs are not directly comparable: the PR run used HF_HUB_OFFLINE=1 (CI-like) while the baseline run had network access, and the baseline misses several models added upstream between the two revisions.

Failure analysis — all failures are environmental, none traced to the PR

1. Offline HF cache misses (the vast majority). The PR run could not download models. Examples: MOSS-Image cache absent → FileNotFoundError: ref-encoder.sock (G1); most G2 failures are the same class; test_omni_processing_correctness explicitly fetches from the Hub (huggingface_hub download errors are imported by the test itself), so it cannot pass offline without a pre-warmed cache.
2. torch teardown crash (G4's 45 ERRORs + UnicodeDecodeError). The crash is in torch/_ops.py::_jit_get_operation from the interpreter-exit weakref teardown of the torch ops cache — it fires after tests complete and pollutes module collection in large combined runs. Reproduced with a standalone run where all 51 tests of tests/helpers/tests/test_assertions.py passed and the crash still occurred at exit. Container torch build issue, unrelated to the PR.
3. Cross-test interference in the big run. The two largest PR-only clusters — test_arg_utils_shared_fields.py::test_aux_outputs_* (12) and test_diffusers_backend.py::test_adapter_guard_* (9) — were re-run individually on dea5597c0: all pass (12 passed / 10 passed).

Conclusion

No PR-introduced regression was found in the local reproduction; every failure cluster is attributable to the offline environment, a torch teardown bug in the container, or run-to-run interference.

[renyongxing]

@hsliuustc0106 The last GPU job of build #17585 has been stuck in waiting for 2+ hours with no log output. The GPU agent appears to be online but shows as Unassigned (not attached to any cluster queue), so it never receives the job. Could someone with admin access re-attach the agent to the GPU queue, then retry the job?