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: @wtomin @RuixiangMa @david6666666 Routing: @wtomin via module of the changed files, module named in the PR description, semantic router, CODEOWNERS; @RuixiangMa via module of the changed files, module named in the PR description, semantic router; @david6666666 via module of the changed files, module named in the PR description, CODEOWNERS @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. |
e6afb18 to
ab9464d
Compare
|
Self-review (requested by the review bot):
Numerics vs the eager chain are not bitwise (≤1 bf16 ulp on ~27% of Q/K elements); the PR body shows this is vLLM |
…ton launch Route FluxAttention (FLUX.1-dev/schnell/Kontext) through the fused joint QKV RMSNorm + RoPE op in the double blocks and the single-stream op in the single blocks; one packed [cos | sin] table per forward via joint_attention_kwargs. Same structure and numerics as Flux.2 (vllm-project#7560). Signed-off-by: HuangYuwei <yuweih205@gmail.com>
… Triton launch Route OvisImageAttention through the fused joint QKV RMSNorm + RoPE op in the double blocks and the single-stream op in the single blocks; one packed [cos | sin] table per forward via joint_attention_kwargs. Same structure and numerics as Flux.2 (vllm-project#7560). Signed-off-by: HuangYuwei <yuweih205@gmail.com>
|
@lishunyang12 Thanks, done: kwargs copied instead of mutated; both kernels call |
Dong1017
left a comment
There was a problem hiding this comment.
Independent verification on our H200 (same stack as your benchmarks: torch 2.13.0+cu130 / vllm 0.28.0 / Triton 3.7.1), tree replicated per the description, all five changed files byte-identical to 63cc63d:
- Unit tests: 19 passed — matches.
- torch.compile(fullgraph) + CUDA graph replay, both ops at production shapes: all outputs bitwise-equal. Worth adding to the test plan.
- FLUX.2-dev (trained weights) e2e, gate off vs 0, E F F E, 3 prompts × 2 seeds, 28 steps @ 1024² — filling the trained-weights gap:
- p50 per image: 13.80 s → 13.42 s (−2.8 %, ≈ −13.6 ms/step); timing stable
- fused arm bitwise-deterministic across independent runs (6/6)
- PSNR eager-vs-fused: min 25.6 / median 35.8 / max 48.6 dB — same composition, background-detail divergence only
Notes: (a) the PSNR floor is lower than klein's ≥37 dB — the 56-block model amplifies the ≤1-ulp perturbation more, worth reflecting; (b) isolated per-block savings (−8.6 %/step) overstate what e2e delivers (−2.8 % ours, −2.9 % klein) — suggest quoting both. And: rope.py is not byte-identical between ff6e906 and main (changed in aff7d64, lazy vllm_flash_attn resolution) — numerics unaffected on 0.28, but the claim needs correcting.
Code findings (details inline): dead table packing on NPU/CPU (agree with @lishunyang12); caller-owned kwargs mutation incl. the pre-existing mask writes; _norm_rope_tile duplication; klein→flux2 private import; untested gate boundary. Also flagging the gate=0 default as a maintainer ratification item.
Integration order — proposal, for discussion with @bugkeep / @WangXukang-cypher (#7417 currently fixes #6982 → #7417 → #7422): #7417 → this PR (rebased onto ops/rope/qk_norm_rope.py) → #7422 (+2 exports). Keeps #7417's AST-verified scope intact, makes the rebase mechanical, resolves the tile duplication here, and leaves #7422's agreed scope intact. This PR first also works but pushes relocation + re-verification onto #7417. Either way, please link RFC #7382 / #7417 / #7422.
…ton launch Route FluxAttention (FLUX.1-dev/schnell/Kontext) through the fused joint QKV RMSNorm + RoPE op in the double blocks and the single-stream op in the single blocks; one packed [cos | sin] table per forward via joint_attention_kwargs. Same structure and numerics as Flux.2 (vllm-project#7560). Signed-off-by: HuangYuwei <yuweih205@gmail.com>
… Triton launch Route OvisImageAttention through the fused joint QKV RMSNorm + RoPE op in the double blocks and the single-stream op in the single blocks; one packed [cos | sin] table per forward via joint_attention_kwargs. Same structure and numerics as Flux.2 (vllm-project#7560). Signed-off-by: HuangYuwei <yuweih205@gmail.com>
|
@Dong1017 Thanks for running it on the trained FLUX.2-dev weights — that is the one arm we could not cover. Folded into the description with credit (per-image −2.8% next to the per-step −8.6%, the 56-block PSNR floor, Last commit: On ordering: either works. The op stays in |
6d30c77 to
55d2328
Compare
… Triton launch Route OvisImageAttention through the fused joint QKV RMSNorm + RoPE op in the double blocks and the single-stream op in the single blocks; one packed [cos | sin] table per forward via joint_attention_kwargs. Same structure and numerics as Flux.2 (vllm-project#7560). Signed-off-by: HuangYuwei <yuweih205@gmail.com>
|
@lishunyang12 Thanks again for the review. Your three points have been addressed: copying Could you please re-review the current head ( Happy to address any remaining concerns. Thanks! |
Merge the updated vllm-project#7560 dependency, preserving this model integration and its numerical and sequence-parallel contracts. Include the shared 64-bit storage-offset fix and regression coverage. Signed-off-by: HuangYuwei <yuweih205@gmail.com>
Merge the updated vllm-project#7560 dependency, preserving this model integration and its numerical and sequence-parallel contracts. Include the shared 64-bit storage-offset fix and regression coverage. Signed-off-by: HuangYuwei <yuweih205@gmail.com>
Merge the updated vllm-project#7560 dependency, preserving this model integration and its numerical and sequence-parallel contracts. Include the shared 64-bit storage-offset fix and regression coverage. Signed-off-by: HuangYuwei <yuweih205@gmail.com>
Merge the updated vllm-project#7560 dependency, preserving this model integration and its numerical and sequence-parallel contracts. Include the shared 64-bit storage-offset fix and regression coverage. Signed-off-by: HuangYuwei <yuweih205@gmail.com>
Merge the updated vllm-project#7560 dependency, preserving this model integration and its numerical and sequence-parallel contracts. Include the shared 64-bit storage-offset fix and regression coverage. Signed-off-by: HuangYuwei <yuweih205@gmail.com>
Omni ReviewBot triage noteResolved as of |
|
The six-PR series is now synced with main ( The synchronization retains main's large-storage-offset fix and tests, carries the same 64-bit row indexing into the joint QKV kernel, and adds four CUDA regression cases covering large offsets in either input stream with both RoPE pairings. The model integrations, floating-point arithmetic, gate default and numerical tolerances are unchanged. Validation: changed-file pre-commit results are recorded as passing; I reran the shared-code/child-integration AST checks and the isolated CPU harness (6 passed, 24 CUDA/Triton cases skipped). That harness uses vLLM import/platform shims; it is not native vLLM or GPU validation. The new GPU regressions and performance measurements have not been rerun on these heads. The previously posted benchmark numbers belong to the earlier measured revisions. New-head CI and maintainer review are still pending. |
Additional numerical/performance experiment (2026-09-21)Keeping the current implementation at Numerical clarification. With vLLM 0.29.0's actual This also corrects the provider attribution in the earlier explanation: the intermediate BF16 rounding before multiplying the norm weight is present in the native implementation, not the v0.29 CUDA implementation. Native and CUDA references must not be conflated; this new diagnosis is scoped to the explicitly checked v0.29 provider, not a reinterpretation of every historical measurement. Avoiding native's extra rounding preserves more intermediate information, but this experiment does not establish that current fusion is more accurate than CUDA against a higher-precision reference. Evidence and setup: one H200, BF16, head dimension 128, epsilon
Performance tradeoff: three arms on identical inputs: original chain / current fusion / exact candidate. No profiler or CUDA Graph replay, median ms per preparation call (not a full attention block or model):
The exact candidate remains 3.58–4.88x faster than the original chain under graph replay, but costs 9.0–10.8% joint / 32.1–33.4% single versus the current fusion. Ordinary uncompiled eager measurements also favor the original PR over the exact candidate; the latter is 2.96–4.89x faster than the original chain. Small eager cases mask part of the kernel penalty behind host launch overhead. The packed RoPE table is reused per forward and excluded from per-block timing; its separate one-time cost was 0.0208–0.0212 ms at B=1 and 0.0331 ms at B=2. No end-to-end speed or image-quality result is inferred from these operator timings. CUDA Graph replay is not evidence of equality against an Inductor-compiled original model, and this standalone matrix does not replace the full checked-in GPU regression suite. Reproduction entry points used in the fixed environment: Given this measured exactness/performance tradeoff, retaining the current implementation, without claiming original-chain bitwise parity. Existing joint-versus-single-fused and same-op compile/graph parity are separate contracts. Follow-up: FP32/FP64 reference error (2026-09-21)We additionally compared the unchanged current fusion, original References lift the same BF16 input/weight values into FP32 or FP64 and perform normalization, weight multiplication and RoPE without BF16 intermediate/output casts. The table below uses the original supplied FP32 cos/sin values for both reference precisions. Error metrics are accumulated in FP64. Current-fused norm values are measured through its real kernel with identity RoPE. The expanded matrix contains 82 single/joint records (41 inputs; 24/48 heads, B=1/2, multiple seeds/weight types, strided inputs and edge cases), aggregating 6,312,665,088 Q/K values per reference/arm. Actual CUDA/native dispatch was asserted. Every arm's repeated outputs were bitwise identical to itself and finite; this is not inter-arm bitwise parity. Element-weighted RMSE, smaller is better:
Thus retaining the current fusion is supported by the measured performance tradeoff and essentially unchanged CUDA-baseline accuracy in this matrix, not a claim that it is substantially more accurate than CUDA. The native comparison demonstrates the cost of its extra intermediate rounding. These are operator results, not pretrained-model quality or universal error bounds. Standalone command: AI assistance: Codex prepared and ran the standalone experiments, checked the raw results, and drafted this update. |
Omni ReviewBot: three questions on the performance claim@yuweih205 this PR reads as a performance or value claim:
Before the full evidence checklist, three short questions:
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. |
Omni ReviewBot: no human activity for 12 days@yuweih205 this pull request has had no human commit, comment or review since 2026-09-21. 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 |
|
@vllm-omni-review-bot Short answers to the three questions:
Base/head SHAs, hardware/software stack, warmup/repeat counts, operator tables, and the higher-precision correctness analysis are already recorded above in the PR. |
Signed-off-by: HuangYuwei <yuweih205@gmail.com>
08dc20e to
80d151c
Compare
Signed-off-by: HuangYuwei <yuweih205@gmail.com>
Signed-off-by: HuangYuwei <yuweih205@gmail.com>
80d151c to
9ed72bf
Compare
|
@lishunyang12 Your three review comments have been addressed on the current 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. |
|
@lishunyang12 One additional numerical point for your re-review: If the concern is that fusion changes floating-point reduction/rounding behavior, operator-level bitwise alignment with a fixed uncompiled eager reference is achievable in principle: we can adjust the kernel's reduction configuration, intermediate precision and rounding points to reproduce that reference. The shared-operator experiment already demonstrated this for the tested vLLM 0.29 CUDA RMSNorm + RoPE contract on H200 (BF16, head dimension 128, epsilon If strict eager operator parity is the preferred acceptance criterion, please let me know; I can adapt and validate that variant for this PR. Equality against an Inductor-compiled full model would require a separate check. |
Signed-off-by: HuangYuwei <yuweih205@gmail.com>
e702a52 to
e0afca0
Compare
Supplemental eager bitwise measurements and queued production checkExisting PR performance results remain as recorded. The retained operator scripts and raw results stay pinned to their original artifact revision. Measured operator case: H200; B1, 512 text + 4096 image tokens, 48 heads, head dimension 128; BF16 activations/weights/RoPE table, full interleaved RoPE, epsilon
Eager CUDA-event medians: six balanced-order samples, 20 warmup calls, 100 calls per sample. These durations cover Q/K(/V) preparation with the packed RoPE table reused; table packing is excluded. They establish operator parity and latency for this case. Prepared production setting: change Actual r3 target: prepared source The three eager arms use the same initialized model/state/inputs:
H200, vLLM Source manifests and reproduction instructions and the eight-worker runner pin the settings. The eight-GPU r3 job is queued in MOVA2.0纯交付分区; full-model parity/performance is not yet established. |
Fresh eager measurements: bitwise reference parity with retained speedup@lishunyang12 The H200 performance follow-up is now complete. The measured CUDA-reference numerical preset has zero output storage-byte mismatches versus the original eager chain, with a measurable transformer-forward saving. Tested implementation: prepared source Settings: Scope: eager, random-weight, full-width shallow FLUX.2 transformer, 48 heads × 128 dimensions, BF16, 4096 image/video tokens and 512 configured text tokens, B1/B2, seed
The original RMSNorm provider is pinned and checked as Timing: H200 141 GB, vLLM Raw records: B1, B2. See the reproduction procedure and all raw results. The fast arm and its separate timing/parity records are retained in those JSON files. AI assistance: Codex ran and checked the GPU measurements and prepared this additive evidence note. Please review the prepared numerical preset and the measured eager cases when convenient. |
Signed-off-by: HuangYuwei <yuweih205@gmail.com>
…K path Flux.2 and Flux.2-klein pass `cuda_rms_norm=True` to the joint and single-stream ops: the fused path is now bit-identical to vLLM RMSNorm + RoPE in eager mode instead of differing by the Triton reduction order. - `fused_joint_qkv_norm_rope` takes `cuda_rms_norm` like the single-stream op (BF16, head_dim = rotary_dim = 128, FMA contraction off) and drops `interleaved`: its half-split mode was never reachable in production. - `pack_qk_norm_rope_table(cuda_rms_norm=True)` returns None outside that geometry, so consumers never fuse outside the bitwise contract. - Shared tests: joint op bitwise against vLLM's CUDA RMSNorm, the Flux.2 eager-chain test exact under `cuda_rms_norm`, compile and CUDA-graph coverage for both modes. Signed-off-by: HuangYuwei <yuweih205@gmail.com>
…nels One-warp programs of eight heads instead of four are 8-15% faster under CUDA-graph replay on H200 across 26 joint and single-stream geometries (320 to 16K tokens, 8 to 64 heads, fast and CUDA-matching modes) with bitwise-identical output. The only slower case is 20 heads in the fast mode (+2.6%), where eight heads pad the head count to 24 slots.
Purpose
Fuse dual-stream Q/K RMSNorm, concatenation and RoPE, and single-stream Q/K preparation, using the shared Triton operators. Integrate Flux.2 and Flux.2-klein.
The fusion paths are selected automatically; there is no fusion environment-variable override. Unsupported devices, dtypes and geometry retain their existing paths. All six PRs contain identical shared operator and test files and only their own model integration; any PR can merge first.
In eager mode the fused path reproduces the vLLM CUDA RMSNorm reduction order and keeps the rotation uncontracted (
cuda_rms_norm=True), so joint Q/K/V and single-stream Q/K are bit-identical to the unfused chain. Compiled graphs (regionaltorch.compile,dynamic=True) keep using the fused custom op. There the unfused chain still runs vLLM's CUDA RMSNorm kernel as an opaque op while Inductor generates the RoPE/concat arithmetic; on the measured shallow shapes the fused output was byte-identical to that compiled chain as well (see Test Result).cuda_rms_norm=Falseat the fused calls restores the previous revision's Triton-order numerics (kernel-level cost of the CUDA-order reduction measured below).Test Plan
7b9b4459d.Test Result
On 2026-10-17 (infra-debug NVIDIA H200; Python 3.12, vLLM 0.31.0, Torch 2.13.0+cu130, CUDA 13.0, Triton 3.7.1; eager, world/TP/SP size 1.)
Shared operator suite: 34 passed, 0 failed/skipped (joint and single-stream ops bit-identical to vLLM's CUDA RMSNorm + eager RoPE under
cuda_rms_norm; compile and CUDA-graph replay bit-identical to the eager op in both modes).Changed-file pre-commit checks (ruff, typos, SPDX, test markers, import checks, mypy) passed.
Eager parity and timing (original arm = table packing disabled, i.e. the unfused chain; bitwise arm = this head's default,
cuda_rms_norm=True; fast arm = the same launch withcuda_rms_norm=False, i.e. the previous revision's Triton-order numerics; CUDA-event medians over three ABBA rounds, each fused arm interleaved with the original arm; exact byte comparison of the full output before and after timing):cuda_rms_norm=False)Kernel alone (joint op, 48 heads, 512 text + 4096 image tokens per sample; CUDA-graph replay medians from the heads-per-program campaign on this operator (26 geometries × fast/bitwise × 4/8/16 heads per program, infra-debug H200, job
hyw-vomni-hp-validation-1017-r1; eight heads were bitwise-identical to four in all 52 path/mode/geometry cells and faster in 51)): four → eight heads per program moves B=1: bitwise 120.8 → 107.1 µs, fast 109.7 → 100.2 µs; the bitwise default is now +6.9% of the fast kernel and −2.4% of the previous revision's fast kernel. At Flux.2-klein's 24 heads: bitwise 62.2 → 55.2 µs (B=1) / 121.2 → 107.4 µs (B=2), fast 56.2 → 51.1 / 109.9 → 100.7 µs.Model level, the bitwise arm's latency is +0.15% (B=1) / −0.19% (B=2) relative to the fast arm — inside the ±0.5% run-to-run spread of this harness. To take the previous numerics instead, pass
cuda_rms_norm=Falseat the fused calls invllm_omni/diffusion/models/flux2/flux2_transformer.pyandflux2_klein/flux2_klein_transformer.py(joint launch, single-stream launch, table packer); this is a code-level switch, no environment variable.Production compile (same harness,
regionally_compile(transformer, dynamic=True)as the diffusion model runner applies it, production IR op priority['vllm_c', 'native']; the previous public head of this PR was run the same way in the same job):Eager in the same run: −2.08% / −2.25%. Both heads compile with the same 22 graph breaks per forward (all from
_vllm_fa2_C.varlen_fwdhaving no fake implementation, which the non-fullgraph production path tolerates), no recompilation across repeated calls, 4 fused launches per forward, and vLLM'srms_norm_kernel(12 launches per forward) present in the unfused chain in both eager and compiled runs.The model-level and request-level rows above were measured at the previous head (
82f8a646d, four heads per program); the eight-heads commit changes only the kernel launch shape with bitwise-identical output, so the parity results carry over and the fused arms gain the kernel delta on top of the figures shown.These are operator, component and forward-level regressions on the stated shapes; pretrained end-to-end accuracy/performance on other models was not rerun for this update.
AI assistance: Claude Code contributed the initial implementation and this revision; Codex simplified automatic selection, implemented and checked CUDA-reference numerics, and ran earlier validation.