Repository navigation
Conversation
|
This PR appears to belong to: docs/design/module/diffusion/diffusion_model_integration.md, docs/design/module/diffusion/index.md. Module owners: @david6666666 @wtomin @Isotr0py Routing: @david6666666 via module of the changed files, module named in the PR description, CODEOWNERS; @wtomin via module of the changed files, module named in the PR description, CODEOWNERS; @Isotr0py via module of the changed files, module named in the PR description @yuweih205, 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. |
ee0501e to
333add5
Compare
|
Self-review: |
Dong1017
left a comment
There was a problem hiding this comment.
Scope. Follow-up to #7560 making Qwen-Image the next consumer of the fused QK-norm-RoPE family: the per-block four-RMSNorm + four-RoPE + three-cat chain becomes one packed launch, with the fp32 [cos | sin] table built once per forward and the eager chain kept as fallback (_FUSED_MIN_TOKENS = 0, env override).
Measured on H200 (nightly diffusers-parity recipe: 512×512, 20 steps, CFG 4.0, seed 42 — complements the PR's fused-vs-eager numerics with vs-diffusers gate numbers):
| path | SSIM | PSNR |
|---|---|---|
| main@21d86ec9 default (fused, fp32 table) | 0.964473 | 28.772 |
| this PR default (packed table) | 0.964473 | 28.772 — trajectory-identical to main |
| this PR, gate off (eager) | 0.964340 | 28.861 — same regime, consistent |
| bf16-coefficient path (reference band, #7513's forced-eager) | 0.985766 | 33.768 |
So the packed refactor preserves diffusers-parity behavior exactly — consistent with the op-level result and the control-arm reasoning in the PR description.
One question this data raises, probably for #7382 rather than here: the "kept in fp32" rationale (identical coefficients across fused/eager) is about internal consistency; vs diffusers, fp32 coefficients diverge on 28.09% of rotated elements (max 0.03125, op-level 1 ulp either way), and the fp32 band (0.9645) sits below #7513's proposed SSIM 0.97 gate while the bf16 band passes with margin (and #7494 shows the same main tree scoring 26.13 on the nightly infra). Since #7513 also rewrites _qwen_image_qk_norm_rope, the regime choice and that threshold seem coupled — might make sense to settle them together as a #7382 contract decision; happy to help with either direction.
333add5 to
0867925
Compare
|
@yuweih205 Please resolve the merge conflicts with the latest main branch |
9f11491 to
f3afda8
Compare
|
@RuixiangMa Done — rebased on latest main and reworked on top of it. main now fuses Q/K RMSNorm + RoPE per stream (#5931) with a BF16 On H200, full 60-block forward: −3.3% / −3.1% / −2.7% eager and −3.3% / −3.3% / −2.3% compiled, at (B=1, 1024²), (B=2, 1024²), (B=2, 512²). In eager the joint output is bitwise identical to the per-stream path (same FP32 coefficients); under |
|
@RuixiangMa One thing worth flagging while you are looking at these: #7560 needs to land first as things stand. It carries the shared op — After #7560 lands, this PR and the other four consumers (#7595 FLUX.1, #7596 HunyuanVideo-1.5, #7597 Z-Image, #7600 Ovis-Image) are independent of each other and can go in any order. I have added the same note to each description. That said, the order is not fixed — tell me which you prefer and I will restructure:
Happy with whichever fits your review flow best. |
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 9 days@yuweih205 this pull request has had no human commit, comment or review since 2026-09-19. 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 routing recordAssigned Strict on zcode (GLM-5.3-Flash) under experiment |
Signed-off-by: HuangYuwei <yuweih205@gmail.com>
2a06195 to
4277824
Compare
Signed-off-by: HuangYuwei <yuweih205@gmail.com>
Signed-off-by: HuangYuwei <yuweih205@gmail.com>
4277824 to
cd2a996
Compare
|
@RuixiangMa The merge conflicts you flagged have been resolved; GitHub currently reports the latest head ( The six-PR series (#7560, #7594, #7595, #7596, #7597, #7600) has also been restructured: each PR contains the same shared operator, tests and documentation changes, plus only its own model integration. There is no longer a prerequisite PR within this series; any one of the six can merge first, once its review and CI requirements are satisfied. This supersedes the earlier comments saying #7560 had to land first. The shared files are identical across the six heads, and local merge simulations passed for all 15 pairs and all 30 ordered pairs with the first PR squash-merged onto main. |
Signed-off-by: HuangYuwei <yuweih205@gmail.com>
Signed-off-by: HuangYuwei <yuweih205@gmail.com>
ebfa21a to
dfc7f36
Compare
Reproducible eager bitwise case: retained full-transformer resultsThe already measured eager configuration preserves the full returned tensor against the historical main per-stream fused baseline, while retaining a 2.72–3.33% transformer-forward saving. The original harness, source overlay and three raw records have now been recovered and made reproducible. These are the September 17 measurements already summarized in this PR; no new GPU measurement is being reported, and the existing compiled and end-to-end performance results are retained. Configuration change: run without the harness's Exact settings:
Timing is CUDA-event p50, A B B A, ten forward calls per segment, with three warmups before each timed segment. Comparison covers the entire returned transformer tensor, converted from BF16 to FP32 before Reproduction uses base PYTHONPATH="$RUN_DIR" "$PYTHON" "$RUN_DIR/qwen_joint_bench.py" --layers 60 --batch 1 --txt 512 --img 4096 --out "$RUN_DIR/b1_4096.json"
PYTHONPATH="$RUN_DIR" "$PYTHON" "$RUN_DIR/qwen_joint_bench.py" --layers 60 --batch 2 --txt 512 --img 4096 --out "$RUN_DIR/b2_4096.json"
PYTHONPATH="$RUN_DIR" "$PYTHON" "$RUN_DIR/qwen_joint_bench.py" --layers 60 --batch 2 --txt 512 --img 1024 --out "$RUN_DIR/b2_1024.json"Raw records: B1 / 4096, B2 / 4096, B2 / 1024. They retain the zero relative-L2 values, exact launch counters and original timings. These establish the stated historical eager case; they are not measurements of the October 8 heads. AI assistance: Codex recovered and checked the existing harness/source hashes/raw records and prepared this additive reproduction note; no GPU case was rerun for this note. |
Fresh eager full-transformer measurement: bitwise joint fusionReran the Qwen-Image eager per-stream-versus-joint case on the current prepared source, with strict storage-byte checks for the complete returned BF16 tensor. Tested source Configuration: full 60-layer random-weight transformer, 24 heads × 128 dimensions, BF16, 512 text + 4096 image tokens, B1/B2, H200 141 GB, world/TP/SP size 1. Constructor seed 0; CUDA generator seed 1 initializes parameters and then inputs. vLLM The reference uses the existing per-stream fused chain plus cats, with only the joint table packer suppressed. Both arms set
Both arms repeat exactly; strict byte parity, including signed zero, is checked before and after timing. Counter wrappers are removed for the measured forwards. CUDA-event medians: three warmups per arm, three ABBA rounds, ten forward calls/sample, six samples/arm. This is a random-weight full-transformer result. B1 raw record, B2 raw record, harness, and the reproduction procedure and all raw results. AI assistance: Codex reran and checked the GPU measurements and prepared this additive evidence note. |
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: 1ef71a57-d4cf-4f07-b058-7cdf689e4956) — 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: aef4a307-2449-4045-aec8-49662a3fbf3a) — 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: 9394ce5d-38cc-4733-bc5f-dc8fd2621749) — check |
vllm-omni-review-bot
left a comment
There was a problem hiding this comment.
Omni ReviewBot review
Changes since the previous review
- 1 new inline finding(s); 1 finding(s) below.
CI at
dfc7f368a121(2026-10-10T10:29:18.793635+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 adds a shared two-stream Triton op that RMSNorms and RoPE-rotates text and image Q/K and gathers V into the joint sequence attention already consumes. Qwen-Image packs one float32 [cos | sin] table per transformer forward and, when that table exists, runs the joint op from every cross-attention block instead of the existing per-stream norm/RoPE path plus three concatenations. Sequence-parallel, non-BF16, and non-CUDA forwards keep the per-stream path.
Change flow
flowchart LR
freqs["[EXISTING] text and image complex freqs"]:::existing
pack["[NEW] packed joint FP32 RoPE table"]:::new
attn["[CHANGED] QwenImageCrossAttention"]:::changed
joint["[NEW] fused_joint_qkv_norm_rope"]:::new
perstream["[EXISTING] per-stream RMSNorm and RoPE"]:::existing
qkv["[CHANGED] joint Q, K, V"]:::changed
freqs --> pack --> attn
attn --> joint --> qkv
attn --> perstream --> qkv
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
Findings
- [P1] Joint gate skips the short-sequence BF16 RoPE path —
vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py:755
Existing thread: #7594 (comment)
Evidence for Joint gate skips the short-sequence BF16 RoPE path
_JOINT_FUSED_MIN_TOKENS is 0, and use_fused_joint launches fused_joint_qkv_norm_rope whenever the packed FP32 table exists. The per-stream helper it replaces still fuses only when B*S >= fused_qk_norm_rope_min_tokens(2048); below that, _qwen_image_qk_norm_rope stays on activation-dtype RotaryEmbedding (torch.real/torch.imag cast to BF16), which is the Diffusers-aligned CUDA path from #7513. Typical text lengths fail that gate (BATCH*512 = 1024 is the existing short-seq test), and a 512² image stream does too, so the default joint path rotates those tokens with the FP32 Triton coefficients instead. test_joint_launch_matches_per_stream_then_cat hides this: force_always_fuse plus use_fused=True compares only with the forced fused kernel. The published bitwise forwards also set VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS=0 on both arms. Enable the joint op only when each stream would already take the fused per-stream kernel, or re-measure the compiled Omni-vs-Diffusers gate (SSIM ≥ 0.97, PSNR ≥ 30) on this default path before treating the FP32 joint launch as the serving default.
🤖 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!
|
|
||
| @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") | ||
| @pytest.mark.parametrize("txt_len,img_len", [(512, 4096), (77, 1024)]) | ||
| def test_joint_launch_matches_per_stream_then_cat(txt_len, img_len, force_always_fuse): |
There was a problem hiding this comment.
[P1] New joint parity tests are invisible to the CUDA model step
Evidence and suggested fix
test_joint_launch_matches_per_stream_then_cat is marked core_model via the module pytestmark and has only a CUDA skipif. It does not carry @pytest.mark.cuda. The ready pipeline step Diffusion · Model Test in .buildkite/cuda/test-ready.yml runs pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model", so those two cases are dropped. The PR's own command is python -m pytest .../test_qwen_image_fused_qk_norm_rope.py -m core_model --run-level=core_model, and the body says adding and cuda would omit them. The new shared-kernel cases live in tests/diffusion/layers/test_fused_qk_norm_rope.py, which has the CUDA marker, but Diffusion · Other Test only selects tests/diffusion/*.py, tests/diffusion/ar_diffusion, and tests/diffusion/layers/test_adalayernorm_*.py, so that file is outside the ready CUDA commands as well. Add @pytest.mark.cuda to the joint Qwen cases so the model step actually runs them.
Omni ReviewBot: finding feedback[p1] Joint gate skips the short-sequence BF16 RoPE path — See the review for details. If you are the PR author and disagree, react 👎 here; the maintainer will see your disagreement. |
Omni ReviewBot: finding feedback[p1] New joint parity tests are invisible to the CUDA model step — See the review for details. If you are the PR author and disagree, react 👎 here; the maintainer will see your disagreement. |
Independent merge update (2026-10-08)
Current head:
cd2a996aon mainc548a110; shared operator commits:bf9acacdand24f9e25e(64-bit indexing). The six PRs #7560, #7594, #7595, #7596, #7597, and #7600 now each contain this identical shared commit plus only their own model integration. Any one of the six can merge first. The other five do not import code from that model.The shared operator, test, and documentation files are byte-identical across all six heads. Local pre-commit checks passed (including ruff, mypy, SPDX, and markdownlint). Git merge simulations found no conflicts for all 15 pairs or for all 30 ordered pairs with the first PR squash-merged onto main. This Mac host has no CUDA/PyTorch test environment, so the GPU results below are historical results on earlier revisions; the new heads await CI and any GPU rerun.
Historical synchronization update (2026-09-21)
Current head:
2a061959c. Merged upstream main1b87115c9ede2d29702c847b6da98f9e854fa593while preserving commit history.At that revision #7560 still landed first. The independent-merge update above supersedes this merge order.
64-bit row addressing to the joint QKV kernel. Added four CUDA regression cases
covering either large-offset stream and both RoPE pairings. Floating-point
arithmetic, model wiring, sequence-parallel behavior, and default gates are unchanged.
ruff, SPDX, test-mark checks, and markdownlint. A Python AST comparison also
verified that the existing model integrations and shared floating-point operations
were preserved (apart from widening address indices).
Operator-level GPU regression update (2026-09-21)
Operator-level correctness and regression tests for the current revision have now passed on real H200 hardware.
tests/diffusion/layers/test_fused_qk_norm_rope.pypassed on [Kernel][Flux.2] Fuse text/image QK RMSNorm + cat + RoPE into one Triton launch #7560 at08dc20e68: 30 passed. The 6 large-storage-offset cases also passed separately, including the 4 new dual-stream cases. Those 6 are included in the 30 and are not counted twice. The shared source and test files are byte-identical across the six current PR heads.2a061959c: 14 passed, covering FP32/FP16 fallback, BF16 and packed-view cases, torch.compile fullgraph, and joint/per-stream parity. These are component/attention-level regressions, not full-model end-to-end tests.Scope: operator-level and model operator-integration regressions are now complete. Pretrained-model end-to-end accuracy and performance ABBA were not rerun. All historical accuracy and performance results below, including the two earlier measurements, are retained with their original revisions and environments; they are not new performance measurements for this revision.
To reproduce the shared operator run, execute from the
tests/directory of #7560 at08dc20e68, with the compatible CUDA/Triton environment and repository test dependencies above. The large-storage-offset cases require slightly more than 4 GiB of free GPU memory:python -m pytest -s -v diffusion/layers/test_fused_qk_norm_rope.py \ -m "core_model and cuda" --run-level=core_model -p no:cacheproviderRun the full operator-integration suite from the
tests/directory of this PR's pinned head:Use
-m core_modelfor the Qwen file without addingand cuda: two joint GPU cases do not carry the CUDA marker and would otherwise be omitted.AI assistance for this synchronization: Codex resolved the merge, updated the indexing
regression coverage, ran the local checks and GPU regressions above, and verified the JUnit results.
Update: against the new main
The joint launch produces the same Q/K/V as main's two per-stream launches followed by the three
cats, so nothing about the images changes — only the launch count and the copies.Numerics (H200, full 60-block transformer, random init, txt 512). Eager: joint vs main's
per-stream path is bitwise identical (
torch.equal, relative L2 exactly 0) at B = 1 and 2,512² and 1024². The packed table carries the same FP32 coefficients
_qwen_image_qk_norm_ropebuilds, so in eager this is equality by construction rather than a tolerance.
Under regional
torch.compilethe two paths are not bitwise equal: relative L2 2.17–2.18% onthe same random-init model, consistently across all three shapes. Only the surroundings differ —
the per-stream arm hands Inductor three
cats to fuse and lay out, the joint arm hands it none —so the graphs either side of the custom op are not the same and the ≤1-ulp consequences amplify
over 60 layers (the same scale as the 1.5–1.6% measured at full depth for FLUX.1 in #7595). We have
not measured the e2e SSIM against Diffusers on the compiled path, so the gate #7513 raised to
0.97 is not something this PR can claim to leave untouched; the eager equality above is what is
established.
Runtime path (asserted in both benchmarks): with the gate off, 120 per-stream launches per
forward (60 blocks × 2 streams) and 0 joint launches; with it on, 60 joint launches and 0
per-stream launches.
Transformer forward, CUDA-event p50, A B B A (3 × 10 after warmup):
Same measurement with the blocks under vLLM-Omni's
regionally_compile(60/60 blocks), which iswhat a default deployment runs:
Peak allocation unchanged (76.2 GiB at B = 2). This is the per-denoising-step cost; the −5.3% per
image reported below was measured against the fully eager pre-#5931 chain and no longer describes
the increment over main. We have not re-run the end-to-end pipeline on the new baseline — with
byte-identical outputs the only change there is this forward saving spread over the denoising loop.
Unit tests: 14 passed on H200 — main's 12 plus two new ones: the joint launch is bitwise equal
to per-stream-plus-cat at (512 + 4096) and (77 + 1024), and the table packer returns
Nonewithoutallocating where the CUDA kernel cannot run.
Sequence parallelism: any SP keeps the per-stream chain.
_sp_planshardsvid_freqswhiletxt_freqsstays replicated, so a joint table built from the two would not line up with the localtokens.
Purpose
Follow-up to #7560 (Flux.2), same optimization for Qwen-Image (
QwenImageCrossAttention, all 60blocks are dual-stream): the per-block chain of four
nn.RMSNorms + four fp32 complex RoPEapplications + three
torch.cats becomes one Triton launch whose outputs are the joint[B, S_txt+S_img, H, D]Q/K/V attention consumes, with no intermediate copies.qk_norm_rope_table; the block forwards it fromjoint_attention_kwargs, and the model packs it once per forward from the complex(txt_freqs, vid_freqs)as an fp32[B*S, D] = [cos | sin]table (text rows first, as theeager
catorders them). Kept in fp32 because the eager path also rotates with unroundedfp32 frequencies (
_apply_qwen_image_rotary_emb).joint_*metadata) and for any dtype/geometry the CUDA kernel cannot take; default gate
_FUSED_MIN_TOKENS = 0,VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENSstill overrides.differ only through fp32 reduction order in
sum(x²)(see numerics).Independent branch: shared operator commit plus this model integration, based directly on main.
Test Plan
vLLM Version: 0.28.0 (torch 2.13.0+cu130, Triton 3.7.1), on the last pre-vLLM-0.29 tree with
qwen_image/andfused_qk_norm_rope.pyfrom main — see #7560 for why.vLLM-Omni Commit: branch on
1b6cd28+ #7560.Hardware: 1× NVIDIA H200.
tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py: table geometry andgate; fused op vs the attention's actual chain (
nn.RMSNorm+_apply_qwen_image_rotary_emb,text-first cat) at (B=1, 77+1024) and (B=2, 512+4096).
Qwen/Qwen-Imagetext-to-image throughOmni(model=..., mode="text-to-image")(
QwenImagePipeline), true CFG 4.0 with a negative prompt, 1024², 50 and 20 steps, 3 prompts ×2 seeds × 2 repeats per arm, gate off vs on, each arm its own process, A B B A. Wall-clock per
generate()including text encoder (Qwen2.5-VL-7B) and VAE.Test Result
Unit tests: 3 passed on H200 (with the op-level suite from #7560: 19 passed).
End-to-end, 1024², CFG (B=2 through the transformer):
Fixed cost outside the denoising loop ≈ 0.14 s; peak allocation unchanged (56.7 GiB).
Numerics:
Each arm is bitwise deterministic across runs.
Op level (unit test): fused vs eager Q/K differ on < 2% of elements, by 1 bf16 ulp — the
fp32 reduction-order residue only, since both sides round once.
Images, eager vs fused on identical (prompt, seed): PSNR median 37.1 dB (24 pairs), minimum
19.1 dB. The low-PSNR pairs are the same scene with small pose/detail differences (e.g. a
horse's leg positions), the usual outcome of a ulp-level perturbation amplified by 20–50 CFG
sampling steps — not a systematic shift.
Control arm. To separate "kernel error" from "any ulp-level change gets amplified", a third
arm runs the eager path with every RMSNorm replaced by a plain torch implementation of the
same single-rounding order (a
sitecustomizepatch in the worker; nothing else changes). Thatarm vs eager: PSNR median 39.9 dB, min 19.1 dB (12 pairs) — the same spread as fused vs eager
(37.1 / 19.1). Fused vs the control arm: 36.1 / 26.9 dB. So a reduction-order-only change
already moves the samples exactly as much as the fused kernel does; there is no additional
error attributable to the kernel.
🤖 Generated with Claude Code
https://claude.ai/code/session_01G8uaASBAbnS8TsFvC5gCUD