[blondeCS]

Contributor intent (required human answers)

These are my own answers. I wrote them in Chinese and translated them with Claude Code.

1. What outcome do you want to achieve?

Let HiDream-O1-Image support CFG parallel flexibly, depending on the case it runs in.

2. What is the highest-leverage variable for that outcome?

How many times each GPU calls the model per denoising step. That is decided at v_guided = self.predict_noise_maybe_with_cfg(...): the loop hands all the arguments to this function in the mixin, and the function decides how many passes to run. With CFG on one GPU it is two calls per step; with CFG on two GPUs, one.

3. How will this work move that variable?

Covered in my answer to question 2.

4. Whose customer or user needs does this address, and what are those needs?

Users who need to run HiDream-O1-Image with CFG parallel on multiple GPUs.

Purpose

Before: HiDreamO1ImagePipeline raises NotImplementedError when cfg_parallel_size > 1, and the feature table lists CFG-Parallel as unsupported for HiDream-O1-Image.

After: the denoising loop goes through CFGParallelMixin. With --cfg-parallel-size 2, one GPU runs the conditional branch and the other the unconditional one, and the two results are exchanged every step. On 2× RTX PRO 6000 Blackwell, a 2048×2048, 50-step image with guidance_scale=5 goes from 21.65 s to 12.07 s (1.79×), and the output is pixel-identical to the single-GPU image. On one GPU, with or without CFG, the loop runs the same operations in the same order as before.

Changes:

• HiDreamO1ImagePipeline inherits CFGParallelMixin, and the cfg_parallel_size > 1 guard is removed.
• New predict_noise: takes the text sample's fields (input_ids, position_ids, attention_mask, full_attn_spans, vinput_mask) and z, t_pixeldit, sigma as keyword arguments, runs the model once, and converts the x-prediction to a velocity as the loop did inline. It takes over the body of _forward_once, whose only caller was that loop.
• The loop calls predict_noise_maybe_with_cfg (with cfg_normalize=False) and scheduler_step_maybe_with_cfg (passing the request's scheduler). The mixin's default combine, n + scale * (p - n), is the expression the loop already used.
• CPU unit tests, a cfg_parallel_size=2 e2e accuracy case next to the existing TP=2 one (they now share a helper), and the feature table cell.
• The SPDX pre-commit hook rewrote the header line of the pipeline file.

Docs: the feature contract docs/design/feature/cfg_parallel.md and the user guide docs/user_guide/diffusion/parallelism/cfg_parallel.md need no change. This PR uses the mixin as the contract documents, and the guide's supported-model list points to the feature table updated here.

Related:

• [RFC]: Continuous Diffusion Model Acceleration Support #1217: I left a note there on Oct 5. HiDream-O1-Image is not in the RFC table; the request comes from the feature table and the guard above.
• [Model] Add HiDream-O1-Image support (text-to-image, editing, personalization, layout-bbox) #5306 by @SOKOUDJOU-LEOPOLD was opened before [Model] Add HiDream-O1-Image support #5194 merged and takes the same route (inherit the mixin, override predict_noise) in its own pipeline. This PR does it for the pipeline that is on main.

Not in this PR

Cache-DiT together with CFG parallel. HiDreamO1TextModel declares has_separate_cfg=True, which does not match one forward per GPU per step under CFG parallel. The open #8998 would resolve that flag to False when the CFG-parallel group has more than one rank, and it covers model-declared adapters like this one, so this PR leaves the flag alone. I have not run Cache-DiT with CFG parallel on HiDream-O1-Image.

Test Plan

vLLM Version: 0.31.0

vLLM-Omni Commit: this branch, rebased on Oct 19 onto 71f0fa4 to resolve a conflict in the feature table. The GPU runs used 45df6b3 on top of 5b95f27. The pipeline file is byte-identical after the rebase, and the commits after it up to 05a39a2 change tests only. Main now includes #8475, which makes the mixin set do_true_cfg on self.transformer when the pipeline has one; this pipeline keeps its model in self.model, so that step does nothing here. I did not rerun the GPU benchmarks after the rebase. The unit tests and pre-commit below were rerun at 05a39a2, the rebased head. On Oct 21, after review, 80c7855 changed predict_noise to take plain tensors instead of HiDreamO1TextSample; at 80c7855 I reran the unit tests, pre-commit and the single-GPU pixel check, but not the two-GPU runs.

Each image run is one fresh process:

python examples/offline_inference/text_to_image/text_to_image.py \
  --model <HiDream-O1-Image> --prompt "A cat is sitting next to a sign that says 'HiDream-O1 vLLM-Omni'" \
  --height 2048 --width 2048 --num-inference-steps 50 --seed 42 \
  --guidance-scale <5.0 or 1.0> [--cfg-parallel-size 2]

The reported time is the script's "Total generation time": the generate call after the engine is up. Model loading is excluded, and there is no warm-up request.

1. Unit tests on CPU: pytest tests/diffusion/models/hidream_o1_image tests/diffusion/distributed/test_cfg_parallel.py. The new file stubs the model call (self.model) and keeps the rest of forward real. One test runs two threads as two CFG ranks, each through its own forward, and checks that rank 0 only runs the conditional branch, rank 1 only the unconditional one, and both end bit-identical to the loop as it was written before.
2. Single GPU, main against this PR (RTX 5090, torch 2.13.0+cu132): guidance 5 and guidance 1, compare pixels and time.
3. Two GPUs (2× RTX PRO 6000 Blackwell, torch 2.13.0+cu130), three runs each: one GPU with guidance 5, one GPU with guidance 1, two GPUs with guidance 5 and --cfg-parallel-size 2.
4. e2e on the same two GPUs: pytest tests/e2e/accuracy/test_hidream_o1_image.py -m "slow and diffusion and cuda and H100 and cards_2" --run-level full_model.
5. pre-commit run on the changed files.

Test Result

Two GPUs (2× RTX PRO 6000 Blackwell, 3 runs each, mean ± std):

Configuration	Time	Peak memory per GPU	Image
1 GPU, guidance 5 (sequential CFG)	21.65 ± 0.01 s	16,387 MiB	reference
1 GPU, guidance 1 (one forward per step)	11.41 ± 0.03 s	16,385 MiB
2 GPUs, guidance 5, cfg_parallel_size=2	12.07 ± 0.10 s	16,771 MiB on each	pixel-identical to the reference

• Speed-up 1.79×. The one-forward-per-step run bounds it at 1.90×; the two-GPU run is 0.66 s per image above that bound. I did not break that 0.66 s down further.
• Repeats of each configuration are pixel-identical to each other.
• Two GPUs with guidance 1 (no CFG to split): 11.49 s, image identical to the single-GPU guidance-1 image (one run).

Single GPU, main against this PR (RTX 5090, one run each):

Guidance	main	this PR	Image
5	35.54 s	35.74 s	pixel-identical
1	18.00 s	17.98 s	pixel-identical

Single GPU, 05a39a2 against 80c7855 (the predict_noise change; RTX 5090, one run each, Oct 21):

Guidance	05a39a2	80c7855	Image
5	35.42 s	35.70 s	pixel-identical
1	17.94 s	17.95 s	pixel-identical

Tests

• Unit (CPU, at 80c7855): 24 passed, 12 skipped. The two extra passes are tests that [Bugfix][Diffusion] Keep TeaCache CFG branches separate in the shared CFG helper #8475 added to test_cfg_parallel.py. The new test file on unmodified main (71f0fa4): 2 failed (test_static_config_accepts_cfg_parallel, test_two_cfg_ranks_run_one_branch_each_and_agree), 2 passed (the single-GPU cases, which run main's own loop against the same stubbed model).
• e2e: 2 passed (tp2 and cfg_parallel2) in 84 s. The H100 mark only selects the tests; they ran on the RTX PRO 6000 pair.
• pre-commit: all hooks pass, on the four changed files at 05a39a2 and on the two files 80c7855 touches.

The numbers from the two machines should not be compared with each other: different GPUs and different torch builds.

Not run

• H100, the hardware the weekly e2e job targets.
• CFG parallel together with Cache-DiT, TeaCache, tensor or sequence parallel, CPU offload, or quantization.
• cfg_parallel_size above 2.
• Online serving; only the offline example and the tests above.

AI assistance: Claude Code helped write the change and the tests, run the benchmarks on my RTX 5090 and on two rented GPUs, translate my answers above, and draft the Purpose, Test Plan and Test Result sections. I reviewed the changes and the results.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	⚠️ Failed 2026-10-17T18:19:12.207378Z	2b31e55	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

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

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/diffusion/index.md, docs/design/module/diffusion/offloader.md, docs/design/module/diffusion/diffusion_model_integration.md.

Module owners: @david6666666 @Isotr0py @princepride

Routing: @david6666666 via module of the changed files, module named in the PR description; @Isotr0py via module of the changed files, module named in the PR description; @princepride via module of the changed files, module named in the PR description

@blondeCS, 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.

[vllm-omni-review-bot]

Omni ReviewBot: three questions on the performance claim

@blondeCS this PR reads as a performance or value claim:

• claim: - Speed-up 1.79×.

Before the full evidence checklist, three short questions:

1. Bottleneck — what is the current bottleneck, and which profile, trace or per-stage measurement shows it?
2. Value — what does the change buy the user or the system (latency, throughput, memory, cost), and at which workload?
3. A/B or ablation — is there a same-workload, same-head/config comparison that isolates each main claim on its own? For stacked optimizations, one number per item rather than a blended delta.

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.

[blondeCS]

Self-review (head 2b31e55, 4 files, +197/−26)

What I checked:

• Scope. The source change is one file, pipeline_hidream_o1_image.py (+28/−21): the class inherits CFGParallelMixin, the cfg_parallel_size > 1 guard is removed, predict_noise is added, and the loop calls the mixin's two helpers. The rest is tests (+168/−4) and one cell in the feature table. No shared code changes.
• One GPU behaves as before. The mixin's sequential path calls predict_noise for the conditional branch, then the unconditional one, and combines them with n + scale * (p - n), the expression the loop used before. On an RTX 5090, main and this PR give pixel-identical images at guidance 5 and at guidance 1.
• Under CFG parallel each GPU runs one branch. Both ranks build both samples; the mixin picks the branch by CFG rank and exchanges the results with one all_gather per step. The two-rank unit test checks that rank 0 only calls the conditional branch, rank 1 only the unconditional one, and both end bit-identical to the loop before this PR. On two GPUs the image is pixel-identical to the one-GPU image.
• Guidance 1 on two GPUs. There is nothing to split, so both ranks run the conditional branch. The image matches the one-GPU guidance-1 image (one run).
• Repository precheck. Before opening, I ran the repo's precheck-pr skill in full mode with Claude Code. One blocking item: the new test faked OmniDiffusionConfig and DiffusionRequestBatch with SimpleNamespace; 2b31e55 uses the real classes. Two warnings: the benchmark command and timed region are now in the description, and I kept the Oct 11 base because main has not changed the pipeline, the mixin or these tests since (the feature table gained an unrelated row, and GitHub reports no conflict).
• Checks. At 2b31e55 on CPU: 22 passed, 12 skipped. The new test file on unmodified main: 2 failed (test_static_config_accepts_cfg_parallel, test_two_cfg_ranks_run_one_branch_each_and_agree), 2 passed. pre-commit passes on the four files, and all seven CI checks pass.

Not checked:

• H100.
• CFG parallel together with Cache-DiT (see [Bugfix] Cache-DiT: no separate-CFG step pairing under CFG parallel #8998), TeaCache, tensor or sequence parallel, CPU offload, or quantization.
• cfg_parallel_size above 2, and online serving.

AI assistance: Claude Code helped draft this comment and ran the precheck. I reran the unit tests and pre-commit myself and reviewed the diff.

[blondeCS]

Answers to the three questions. All numbers come from the runs in the description. Two GPUs: 2× RTX PRO 6000 Blackwell Server Edition, driver 580.82.09, torch 2.13.0+cu130, vLLM 0.31.0, HiDream-ai/HiDream-O1-Image at revision 0b0901d9. Workload: one 2048×2048 image, 50 steps, seed 42, batch 1, one request per fresh process, no warm-up, 3 runs per configuration, mean ± std. Base 5b95f27; the GPU runs used head 45df6b3, and the later commits change tests only.

1. Bottleneck. With CFG on, each of the 50 steps runs the transformer twice, one call after the other, on one GPU. The timings show the transformer calls dominate: guidance 5 (100 calls) takes 21.65 ± 0.01 s, guidance 1 (50 calls) 11.41 ± 0.03 s. The difference, 10.24 s for 50 calls, is about 0.20 s per call, which leaves about 1.2 s of the guidance-1 run outside the transformer calls. This is arithmetic on end-to-end timings, not a profile; I have no profiler trace.

2. Value. Latency of one CFG image for users with two GPUs: 21.65 s → 12.07 ± 0.10 s (1.79×) at 2048², 50 steps, guidance 5, with a pixel-identical image. Before this PR, cfg_parallel_size=2 raised NotImplementedError for this model. The cost is a second GPU holding a full copy of the model, plus 384 MiB more peak memory per GPU (16,387 → 16,771 MiB). It is a latency feature, not a throughput one: per image it uses 2 × 12.07 = 24.1 GPU-seconds against 21.65 on one GPU.

3. A/B. There is one change, so one comparison: one GPU against two GPUs with --cfg-parallel-size 2, on the same machine, head, workload and seed. 21.65 ± 0.01 s against 12.07 ± 0.10 s.

• The one-call-per-step run (guidance 1, one GPU, 11.41 ± 0.03 s) is the floor for two GPUs and bounds the speed-up at 1.90×. The two-GPU run is 0.66 s above that floor; I did not break that down further.
• Quality: the two-GPU image is pixel-identical to the one-GPU image (max abs diff 0), and the repeats of each configuration are pixel-identical to each other.
• One GPU, main against this PR (RTX 5090, torch 2.13.0+cu132, one run each): 35.54 s against 35.74 s at guidance 5, 18.00 s against 17.98 s at guidance 1, pixel-identical in both. With one run each there is no spread, so these show that nothing changed much, not an exact difference.

AI assistance: Claude Code helped draft this reply from my benchmark logs. I checked the numbers against the logs.

[tzhouam]

double check this sample data type, thanks

[blondeCS]

Thanks, changed in 80c7855. predict_noise no longer takes the model-specific HiDreamO1TextSample. It takes the sample's fields as keyword arguments (input_ids, position_ids, attention_mask, full_attn_spans, vinput_mask) plus z, t_pixeldit and sigma, and the loop unpacks each sample into the kwargs it hands to predict_noise_maybe_with_cfg. full_attn_spans stays a list of index spans because that is what the model takes. _forward_once had no other caller, so its body moved into predict_noise unchanged.

@hsliuustc0106 this is also the point you raised in today's meeting.

Checked at 80c7855:

• CPU unit tests: 24 passed, 12 skipped. The CFG-Parallel test now stubs the model call instead of _forward_once, so the vinput_mask selection and the velocity conversion run in the test.
• pre-commit on the two changed files passes.
• RTX 5090, 2048×2048, 50 steps, seed 42: the images at guidance 5 and 1 are pixel-identical to 05a39a2.

I did not rerun the two-GPU benchmark. The description is updated.

AI assistance: Claude Code helped make this change, run the checks and draft this reply. I reviewed the change and the results.

[blondeCS]

@tzhouam updated in 80c7855, could you take another look when you have time? Thanks!

[vllm-omni-review-bot]

Omni ReviewBot: no human activity for 7 days

@blondeCS this pull request has had no human commit, comment or review since 2026-10-27. Please confirm the current plan and next step. 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.

[blondeCS]

Nothing pending on my side: the only review comment was addressed in 80c7855 on Oct 21, and a re-review is requested.