Repository navigation
[Diffusion] [batch invariance] Install batch invariance in the diffusion worker and document its verified scope [2/2] - #5935
Conversation
…fusion Batch-invariant diffusion needs a portable RNG identity: the same request must reproduce bit-identically regardless of how it was batched. A pre-seeded torch.Generator cannot provide that, because its state is consumed in place and does not survive the cross-process hop to a diffusion stage worker. Add a request-level contract that, when batch invariance is on, requires exactly one explicit integer seed within torch.Generator.manual_seed range and rejects generator input. The switch is tri-state: VLLM_OMNI_DIFFUSION_BATCH_INVARIANT unset follows vLLM's global VLLM_BATCH_INVARIANT, so with both unset behavior is byte-for-byte unchanged; setting it explicitly overrides the global in either direction, letting a mixed AR + diffusion pipeline enable invariance for its LLM stage without constraining the diffusion stage. StagePool now normalizes diffusion params through one helper on both the initial and streaming admission paths, and copies params per dispatch so a stage that initializes a generator in place cannot leak that state into the next dispatch. Internal warmup/dummy runs keep their existing internal seed fallback. Signed-off-by: CrazyCoder <18235787078@163.com>
Conflict: vllm_omni/engine/stage_pool.py only, against upstream 1c6e731 "fix(diffusion): convert caller sampling params instead of discarding" (vllm-project#4115). Resolution: drop our hand-rolled conversion `OmniDiffusionSamplingParams(seed=getattr(params, "seed", None) if enabled else None)`. It kept only the seed and discarded every other caller-supplied field, which is exactly the bug vllm-project#4115 fixed. Upstream's `OmniDiffusionSamplingParams.from_params()` covers the same responsibility and does it completely, so `_prepare_diffusion_params` now delegates normalization to it. Kept from this branch: - the batch-invariant seed contract check (`validate_batch_invariant_diffusion_seed`) - `copy.copy(params)` when the feature is enabled, so a stage initializing the generator in place cannot leak that state into the next dispatch - the lazy `vllm_omni.diffusion.batch_invariance` import `diffusion_batch_invariant_enabled()` stays ahead of `from_params()` so that when the env value is invalid and the caller's params type is unsupported, the exception type raised first is unchanged. Test update: test_feature_off_plain_sampling_params_preserve_legacy_seed_drop asserted the seed-drop behavior vllm-project#4115 removed. Renamed to test_feature_off_plain_sampling_params_keep_caller_seed and its assertion now expects the caller's seed=17 to survive. The other five tests still hold as-is.
…fusion Batch-invariant diffusion needs a portable RNG identity: the same request must reproduce bit-identically regardless of how it was batched. A pre-seeded torch.Generator cannot provide that, because its state is consumed in place and does not survive the cross-process hop to a diffusion stage worker. Add a request-level contract that, when batch invariance is on, requires exactly one explicit integer seed within torch.Generator.manual_seed range and rejects generator input. The switch is tri-state: VLLM_OMNI_DIFFUSION_BATCH_INVARIANT unset follows vLLM's global VLLM_BATCH_INVARIANT, so with both unset behavior is byte-for-byte unchanged; setting it explicitly overrides the global in either direction, letting a mixed AR + diffusion pipeline enable invariance for its LLM stage without constraining the diffusion stage. StagePool now normalizes diffusion params through one helper on both the initial and streaming admission paths, and copies params per dispatch so a stage that initializes a generator in place cannot leak that state into the next dispatch. Internal warmup/dummy runs keep their existing internal seed fallback. Signed-off-by: CrazyCoder <18235787078@163.com>
…ion worker and document its verified scope [2/2] Layer 2 of the three-layer batch-invariance design (vllm-project#5512). This stacks on vllm-project#5815 (the seed contract) and should be reviewed after it. The seed contract makes diffusion randomness reproducible but never installs vLLM's batch-invariant operator overrides on the diffusion side, so a caller who supplies a seed still runs through kernels whose reduction partitioning depends on batch composition. The low bits still move. This commit installs them, extends the request contract from "a seed exists" to "the seed alone decides the RNG identity", and states the verified scope in documentation rather than enforcing it. Worker bootstrap (diffusion/worker/diffusion_worker.py): resolve the switch during startup, after the CUDA device is selected and before init_distributed_environment, so vLLM's NCCL determinism overrides land before the communicator exists. Returns silently on ROCm and non-CUDA devices rather than raising. Request contract (diffusion/batch_invariance.py, diffusion/request.py): validate_batch_invariant_diffusion_request additionally rejects generator_device, latents and sigmas. These are not recipe constraints but the premise of "same seed => same output" itself -- caller-supplied latents mean the seed no longer decides the initial noise, and sigmas rewrite the schedule the seeded trajectory follows. The call site is OmniDiffusionRequest.__post_init__, which sees every construction path including callers that build a request directly; warmup stays exempt via is_dummy_run(). Scope is documented, not gated (docs/features/batch_invariance.md). Configurations outside the measured tuple run rather than being refused: determinism holds for the operators vLLM actually replaces, and everything else is unverified rather than unsupported. This follows vLLM's own approach, which lists validated models and notes that others may also work. The page records the measured SD3 single-GPU evidence, why video and audio pipelines have no evidence (upstream overrides no convolution operator), that multi-GPU is unverified in both directions along with the specific audit points it would need, and one negative result: SM120 (RTX 5090) has been observed to lose batch invariance. It also records that the seed contract is not device-aware, so on ROCm and non-CUDA devices an enabled switch still constrains requests while the bootstrap installs nothing -- VLLM_OMNI_DIFFUSION_BATCH_INVARIANT=0 is the way out there. Tests pin each retained request check, that an off-table configuration is accepted so the removed gate cannot return unnoticed, that the seed explicitness record is taken before the random-seed fallback, and that the worker bootstrap is a no-op when the switch is off. Verified on CPU only. The SD3 bit-identity measurements quoted in the docs were collected before the GPU harness was split out of this branch; no GPU run backs this revision. Signed-off-by: CrazyCoder <18235787078@163.com>
hsliuustc0106
added
diffusion
|
This PR touches vllm_omni/diffusion/, tests/diffusion/, docs/, docs/features/, tests/engine/ (9 files). Based on CODEOWNERS coverage of the changed files, the most-related reviewers appear to be: Could one of you take a look when you get a chance? Thanks! |
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
|
This PR appears to belong to: docs/design/module/diffusion/index.md. Module owners: @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment. |
Omni ReviewBot triage noteAutomated triage of commit
These are automated triage suggestions only — the final decision belongs to the maintainers. |
Omni ReviewBot: no human activity for 8 days@Ronnie-Rui 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. |
Signed-off-by: CrazyCoder <18235787078@163.com>
Signed-off-by: CrazyCoder <18235787078@163.com>
| # one. Internal warmup requests are exempt: they are constructed by the | ||
| # engine itself and legitimately rely on that fallback. | ||
| if not self.is_dummy_run(): | ||
| validate_batch_invariant_diffusion_request(self) |
There was a problem hiding this comment.
Validate external requests only; Helios rebuilds requests after the runner initializes their generators.
There was a problem hiding this comment.
Thanks @lishunyang12! Moved validation to request admission so internal requests can be rebuilt after generator initialization. Also preserved explicit-seed checks.
Omni ReviewBot: no human activity for 7 days@Ronnie-Rui this pull request has had no human commit, comment or review since 2026-09-17. 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. |
Omni ReviewBot: no human activity for 14 days@Ronnie-Rui this pull request has had no human commit, comment or review since 2026-09-17. 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. |
Omni ReviewBot routing recordAssigned Strict on zcode (GLM-5.3-Flash) under experiment |
Omni ReviewBot attempt recordReview attempt ended as failed (step 'review' (agent.review_diff): unhandled error: RuntimeError: zcode exited 1 without a result event: statusCode: undefined } Error: Turn execution failed (traceId: e97c615c-b1d6-4bac-829a-227916fa882f) — check |
Omni ReviewBot attempt recordReview attempt ended as failed (step 'review' (agent.review_diff): unhandled error: RuntimeError: zcode exited 1 without a result event: statusCode: undefined } Error: Turn execution failed (traceId: 91e11832-80a4-40a6-81a0-45217021cc60) — check |
Omni ReviewBot attempt recordReview attempt ended as failed (step 'review' (agent.review_diff): unhandled error: RuntimeError: zcode exited 1 without a result event: statusCode: undefined } Error: Turn execution failed (traceId: 039c7a27-c9bb-41b1-9043-dfd21dcb88fa) — check |
vllm-omni-review-bot
left a comment
vllm-omni-review-bot
left a comment
There was a problem hiding this comment.
Omni ReviewBot review
Changes since the previous review
- 2 new inline finding(s); 0 finding(s) below.
CI at
4d8ddde1ea45(2026-10-09T13:46:50.521443+00:00): required check(s) blocking:buildkite/vllm-omni(missing).
Note: The assigned review arm
strict/zcode/GLM-5.3-Flashcould not complete this review, so it was produced by the fallback armdirect/cursor/auto. It is excluded from the routing experiment.
Full review analysis
PR description
This change turns vLLM's batch-invariant kernels on inside the diffusion worker and ties that mode to a seed-only RNG contract. After the CUDA device is selected and before init_distributed_environment, the worker calls init_batch_invariance() when the diffusion switch is on, and it skips ROCm and non-CUDA devices without installing overrides. Non-warmup OmniDiffusionRequest construction then requires an explicit integer seed and rejects generator, generator_device, latents, and sigmas. Stage-pool admission checks only the seed and generator, and copies sampling params so a later in-place generator cannot leak into the next dispatch. The new feature page records the measured SD3 single-GPU scope and leaves other recipes unverified rather than refused.
Change flow
flowchart LR
Env["[EXISTING] VLLM_BATCH_INVARIANT"]:::existing
Switch["[NEW] Diffusion batch-invariance switch"]:::new
Worker["[CHANGED] DiffusionWorker.init_device"]:::changed
Request["[CHANGED] OmniDiffusionRequest.__post_init__"]:::changed
Pool["[CHANGED] StagePool diffusion admission"]:::changed
Docs["[NEW] Batch invariance feature doc"]:::new
Env --> Switch
Switch --> Worker
Switch --> Request
Switch --> Pool
Switch --> Docs
classDef existing fill:#e5e7eb,stroke:#6b7280,color:#111827
classDef changed fill:#fef3c7,stroke:#d97706,color:#451a03,stroke-width:2px
classDef new fill:#dcfce7,stroke:#16a34a,color:#052e16,stroke-width:2px
classDef removed fill:#fee2e2,stroke:#dc2626,color:#450a0a,stroke-width:2px
See inline comments below.
🤖 This review was generated by InferMatrix Copilot, an open-source repo-maintenance agent for PR review, CI debugging and issue triage. Try it on your own repo, and ⭐ star it if it helped!
Omni ReviewBot: no human activity for 30 days@Ronnie-Rui this pull request has had no human commit, comment or review since 2026-09-17. Please consider closing this PR if the work is no longer planned. 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. |
Merge current main and apply the reviewed seed contract fixes. Validate image seeds before fallback and use request-scoped client errors. AI assistance: Codex assisted with implementation and local validation. Signed-off-by: CrazyCoder <18235787078@163.com>
Merge the updated seed contract and validate external RNG inputs before replica binding and preprocessing while preserving internal reconstruction. Use the current vLLM bootstrap import and retain OpenPI seed provenance. AI assistance: Codex assisted with implementation and local validation. Signed-off-by: CrazyCoder <18235787078@163.com>
…pipelines On one L4 with the tiny pipelines in eager mode, the default kernels are not batch invariant: the co-batched Qwen-Image output differs from the solo output in about 40 percent of uint8 elements (max abs diff 4) and the LTX2 video in 0.75 percent (max abs diff 2). That is the feature-off control vllm-project#5512 describes, before vLLM's batch-invariant overrides are installed in the diffusion worker (vllm-project#5935). Add batch_invariance_xfail_reason so a measured pipeline keeps running the check as a strict xfail and flips the expectation when it passes. Run the tasks without subtests so the mark applies to the whole case, print the dispatched batch sizes for the log, and opt in QwenImagePipeline and LTX2Pipeline. FluxPipeline is no longer opted in: its tiny builder fetches gated configs, so it was not measured. Signed-off-by: Rishabh Sinha <rsinha17@terpmail.umd.edu>
rishabhsinha17
left a comment
rishabhsinha17
left a comment
There was a problem hiding this comment.
Installs as described, with a measured limit below. The project needs it because Part 1's seed contract runs through batch-dependent kernels until the worker installs vLLM's overrides.
Verified at 1c6fb99. Worker bootstrap runs after device selection and before distributed init (diffusion_worker.py:397 and :416), aligns VLLM_BATCH_INVARIANT for vLLM's own check, and checks no capability. Admission validates before binding (stage_pool.py:1074).
Measured on one L4 (SM 8.9, vllm 0.32.0, tiny random-weight pipelines, co-batching confirmed, head merged with main, the base predates #9066). Switch off, Qwen-Image differs in 310788 of 786432 elements and LTX2 in 51179 of 7077888. Switch on (worker mode flag True), Qwen-Image still differs in 321413 and 317207 and LTX2 in 54491 and 3096309 over two runs. The overrides leave these pipelines unchanged.
Layer 2 of #5512 became documentation rather than refusal, a call for @SamitHuang and @hsliuustc0106.
AI assistance: Claude Code drafted this review and ran the checks named above.
|
|
||
| Every measurement above is single-GPU, with all 12 parallelism degrees in `DiffusionParallelConfig` at 1. Multi-GPU diffusion batch invariance is therefore **unverified in both directions** — no evidence that it holds, none that it fails. It is not gated, so it runs. That is a structural gap rather than a sampling one: on one GPU the diffusion stage performs no collective communication at all, so the single-GPU runs never exercise the code multi-GPU determinism would rest on. | ||
|
|
||
| An omni-specific point needs auditing before trusting that path. vLLM sets `disable_custom_all_reduce = True` whenever `VLLM_BATCH_INVARIANT` is on, in `vllm/config/parallel.py`. That name appears zero times anywhere in `vllm_omni`, and diffusion calls `torch.distributed.all_reduce` directly rather than routing through vLLM's `CustomAllreduce`, so the upstream protection does not obviously carry over. Two upstream issues describe the failure modes to expect: vllm#50136 (custom all-reduce eligibility depends on a size threshold, so kernel selection becomes a function of batch composition) and vllm#30321 (DP + EP inconsistency). |
There was a problem hiding this comment.
This sentence is stale at vllm 0.32.0, the version the measurement above ran on. vllm/config/parallel.py no longer disables the custom all-reduce for batch invariance. override_envs_for_invariance() in vllm/model_executor/determinism/batch_invariant.py pins VLLM_CUSTOM_ALLREDUCE_ALGO=1stage and VLLM_ALLREDUCE_USE_SYMM_MEM=0 instead, so the multi-GPU audit should look for those two variables and the one-stage kernel rather than for disable_custom_all_reduce.
|
|
||
| The worker does not check compute capability. Any other CUDA GPU is **unverified, not rejected**: the switch is honoured, the engine runs, and determinism holds only for the operators vLLM actually replaces. Do not read that as "newer is safer" — vLLM *branches* on capability rather than requiring a floor, so a result in one capability bracket says nothing about another. | ||
|
|
||
| One negative result is already known: **on SM120 (compute capability 12.0, RTX 5090) batch invariance does not hold.** That is an observed failure, not a prediction — identical requests diverge across batch sizes with the switch on. |
There was a problem hiding this comment.
Is there a link or a recorded run for this result? docs/contributing/ci/lessons.md L002 asks that a measured result name its hardware plus the revision and command that produced it, and this is the only measured negative on the page.
| if before is sentinel: | ||
| os.environ.pop("VLLM_BATCH_INVARIANT", None) | ||
| else: | ||
| os.environ["VLLM_BATCH_INVARIANT"] = before |
There was a problem hiding this comment.
mypy-3.10 flags this line because before carries type object through the sentinel. CI skips that hook, but a local pre-commit run fails on it. Environment values are always strings, so os.environ.get with None as the unset marker keeps the behaviour and satisfies mypy.
Purpose
Design issue #5512 splits batch-invariance enforcement into three layers. This PR is Layer 2, and it stacks on #5815.
#5815 makes diffusion randomness reproducible: a request must carry one explicit integer seed. What it does not do is install vLLM's batch-invariant operator overrides on the diffusion side, so the same seed still runs through kernels whose reduction partitioning depends on batch composition, and the low bits still move. This PR closes that half.
Scope is communicated through documentation rather than refusal. determinism holds for the operators vLLM actually replaces, and anything else is unverified rather than unsupported. This follows vLLM's own approach of listing validated models while noting that others may also work.
What this PR changes
Production:
vllm_omni/diffusion/worker/diffusion_worker.py— bootstrap batch invariance during worker startup, after the CUDA device is selected and beforeinit_distributed_environment, so vLLM's NCCL determinism overrides land before the communicator exists. Returns silently on ROCm and non-CUDA.vllm_omni/diffusion/batch_invariance.py— addvalidate_batch_invariant_diffusion_request, which additionally rejectsgenerator_device,latentsandsigmas. These are not recipe constraints: caller-suppliedlatentsmean the seed no longer decides the initial noise, andsigmasrewrite the schedule the seeded trajectory follows.vllm_omni/diffusion/request.py— record whether the seed was caller-supplied, and run the request contract from__post_init__, which sees every construction path. Warmup requests stay exempt viais_dummy_run().Docs:
docs/features/batch_invariance.md(new) anddocs/.nav.yml— enabling, the seed contract, the measured evidence table, why only SD3 has evidence, the multi-GPU position, and known limitations.Test Plan
pytest tests/diffusion/test_batch_invariance_seed_contract.py tests/diffusion/test_diffusion_request.py, plus the Layer 1 suites from #5815 (tests/engine/test_stage_pool_diffusion_seed.py).Measured batch-invariance evidence
Shared recipe:
StableDiffusion3Pipeline(stable-diffusion-3.5-medium), 512x512, 8 steps,torch.bfloat16,TORCH_SDPA,enforce_eager=True,num_gpus=1. Four requests A/B/C/D each pin their own prompt and seed (A ="a red ceramic cup on a plain white table", seed 101); the BS4 batch is submitted as[B, A, D, C]and only the A path is compared. Every hash is a per-sample tensor sha256 overshape=[1, 3, 512, 512],nbytes=1572864, bf16 — not a whole-batch tensor.RTX 4090 / SM 8.9 — invariance holds
a731ab61eac277b30a2fcbaf547e812251fd184b605360d8d453603c6f79b478a731ab61eac277b30a2fcbaf547e812251fd184b605360d8d453603c6f79b478daedf31688ccd792404b956873aa82441be673efab8ccd54b7a5511d0be72db6b3a971fe2ba4be0f7d6cc659600560a9236f40fc0ebdfe1b098a899dfcf4ff40Rendered side by side, with the per-pixel absolute difference on the right. Both heatmaps use the same color scale (
0→0.1934), so their shading is directly comparable.Gate on — the difference image is empty because every element is exactly equal:
[B, A, D, C]Gate off — same prompt, same seed, same recipe, batch composition the only variable:
[B, A, D, C]The gate-on tensor is not degenerate:
std = 0.4988across 1,974 distinct sampled values, so "identical" is not two constant images agreeing.Sampling self-check: both groups enumerated all four paths of their BS4 run and got
distinct=4/4, with the A path matching the hashes above character for character, so the comparison is not reading the wrong sample. At the operator level, inside the same worker process, one row multiplied alone versus the same row multiplied inside a batch gavemax_abs_diff=1.0(not equal) beforeinit_batch_invariance()and0.0(equal) after. Environment: RTX 4090 24564 MiB,get_device_capability(0) = [8, 9],is_device_capability_family(80) = True(both printed at runtime, not inferred), driver 580.105.08, torch 2.11.0+cu130, vllm 0.25.0.Not covered yet: multi-GPU and video
No experiment was run for either. This is a scope statement, not a measured negative result.
Multi-GPU — every run above used
num_gpus=1, so no collective communication was exercised. Upstream's own batch-invariance CI also runs at TP=1 only, and batch invariance at TP>1 is still an open upstream bug (vllm#51290). A run today would most likely report "not invariant" without being able to attribute it, so multi-GPU stays unverified in both directions.Video — upstream's overrides cover matmul-shaped operators plus softmax and
mean.dim; there is no batch-invariant convolution orgroup_normoverride at all. Video pipelines are built onConv3d, where the time dimension enlarges the reduction and pushes implementations toward shape-dependent tiling. A convolution is not disqualifying by itself — SD3's verified path already contains aConv2dpatch embed whose reduction extent is independent of batch composition — but the 3D case can only be settled once upstream has batch-invariant convolution kernels.