Skip to content

[Kernel][Flux.2] Fuse text/image QK RMSNorm + cat + RoPE into one Triton launch - #7560

Open
yuweih205 wants to merge 7 commits into
vllm-project:mainfrom
yuweih205:perf/flux2-fused-qk-norm-rope
Open

yuweih205 wants to merge 7 commits into
vllm-project:mainfrom
yuweih205:perf/flux2-fused-qk-norm-rope

Conversation

@yuweih205

@yuweih205 yuweih205 commented Sep 15, 2026 •

Copy link
Copy Markdown
Contributor

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 (regional torch.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=False at the fused calls restores the previous revision's Triton-order numerics (kernel-level cost of the CUDA-order reduction measured below).

Test Plan

  • Current commit: 7b9b4459d.
  • infra-debug NVIDIA H200; Python 3.12.13, vLLM 0.31.0, Torch 2.13.0+cu130, CUDA 13.0, Triton 3.7.1.
  • Repository pytest configuration and CI markers enabled.
python -m pytest -s -v tests/diffusion/layers/test_fused_qk_norm_rope.py \
  -m core_model --run-level core_model

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 with cuda_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):

Case Original Bitwise (default) Fast (cuda_rms_norm=False) Bitwise vs fast Bytes differing vs original: bitwise / fast
Flux.2 shallow forward, B=1 38.78 ms 38.04 ms (−1.89%) 37.99 ms (−3.18%) +0.15% 0 / 434932 of 1048576 (max abs 0.0781)
Flux.2 shallow forward, B=2 77.89 ms 75.49 ms (−3.09%) 75.64 ms (−2.89%) −0.19% 0 / 870638 of 2097152 (max abs 0.0781)

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=False at the fused calls in vllm_omni/diffusion/models/flux2/flux2_transformer.py and flux2_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):

Case Original (compiled) Fused (compiled) Fused vs original output
Flux.2 shallow forward, compiled, B=1 37.01 ms 36.35 ms (−1.77%) 0 bytes differing
Flux.2 shallow forward, compiled, B=2 74.60 ms 73.16 ms (−1.93%) 0 bytes differing

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_fwd having no fake implementation, which the non-fullgraph production path tolerates), no recompilation across repeated calls, 4 fused launches per forward, and vLLM's rms_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.

@vllm-omni-review-bot

Copy link
Copy Markdown

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.

@yuweih205
yuweih205 force-pushed the perf/flux2-fused-qk-norm-rope branch from e6afb18 to ab9464d Compare September 15, 2026 04:12
@yuweih205

Copy link
Copy Markdown
Contributor Author

Self-review (requested by the review bot):

  • vllm_omni/diffusion/layers/fused_qk_norm_rope.py: factored the interleaved/half-split Q/K RMSNorm+RoPE arithmetic of _qk_norm_rope_kernel into _norm_rope_tile (no arithmetic change; the existing single-stream kernel and its tests are untouched) and added _joint_qkv_norm_rope_kernel + fused_joint_qkv_norm_rope (custom op vllm_omni.fused_joint_qkv_norm_rope, fake impl, eager fallback). The joint kernel is bitwise equal to running the single-stream kernel per stream.
  • vllm_omni/diffusion/models/flux2/flux2_transformer.py: Flux2Attention (double blocks) takes the joint op, Flux2ParallelSelfAttention (single blocks) takes the existing fused_qk_norm_rope(interleaved=True); one packed rope table per forward through joint_attention_kwargs. SP paths and any unsupported dtype/geometry keep the original eager chain, which is bitwise identical to main (checked at model level with the gate off). Default gate _FUSED_MIN_TOKENS = 0 for Flux.2, still overridable through VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS.
  • tests/diffusion/layers/test_fused_qk_norm_rope.py: joint op vs per-stream single kernel (bitwise, incl. empty streams, strided chunked-QKV views, fp32/bf16 tables), vs eager reference (both pairings), vs Flux.2's real chain (vLLM RMSNorm + RotaryEmbedding), and shape validation.
  • docs/configuration/environment_variables.md: Flux.2 default noted in the env var table.

Numerics vs the eager chain are not bitwise (≤1 bf16 ulp on ~27% of Q/K elements); the PR body shows this is vLLM rms_norm's double rounding vs the kernel's single rounding, with the kernel closer to an fp32 reference and its own accumulation-order contribution at 0.00%.

Comment thread vllm_omni/diffusion/models/flux2/flux2_transformer.py Outdated
Comment thread vllm_omni/diffusion/layers/fused_qk_norm_rope.py Outdated
Comment thread vllm_omni/diffusion/models/flux2/flux2_transformer.py Outdated
@hsliuustc0106 hsliuustc0106 added the high priority high priority issue, needs to be done asap label Sep 16, 2026
yuweih205 added a commit to yuweih205/vllm-omni that referenced this pull request Sep 16, 2026
…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>
yuweih205 added a commit to yuweih205/vllm-omni that referenced this pull request Sep 16, 2026
… 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>
@yuweih205

yuweih205 commented Sep 16, 2026 •

Copy link
Copy Markdown
Contributor Author

@lishunyang12 Thanks, done: kwargs copied instead of mutated; both kernels call _norm_rope_tile; pack_qk_norm_rope_table returns None (no allocation) unless the CUDA kernel would run. CUDA path unchanged (bitwise identical, same timings).

@Dong1017 Dong1017 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread vllm_omni/diffusion/models/flux2/flux2_transformer.py Outdated
Comment thread vllm_omni/diffusion/models/flux2/flux2_transformer.py Outdated
Comment thread vllm_omni/diffusion/layers/fused_qk_norm_rope.py Outdated
Comment thread vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py Outdated
Comment thread vllm_omni/diffusion/layers/fused_qk_norm_rope.py
yuweih205 added a commit to yuweih205/vllm-omni that referenced this pull request Sep 17, 2026
…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>
yuweih205 added a commit to yuweih205/vllm-omni that referenced this pull request Sep 17, 2026
… 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>
@yuweih205

Copy link
Copy Markdown
Contributor Author

@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, rope.py wording corrected).

Last commit: QK_NORM_ROPE_TABLE_KEY moved into layers/ (no more klein→flux2 import), SP skip moved inside pack_qk_norm_rope_table, CPU tests for the gate boundary / env override / SP skip, and a CUDA test that both ops are bitwise-equal to eager under torch.compile(fullgraph=True) and CUDA-graph replay. The packing gate, the kwargs copy (the mask writes now land on the copy too) and the _norm_rope_tile consolidation were already in 6dd5d04.

On ordering: either works. The op stays in layers/fused_qk_norm_rope.py, so we would rather land on the current path and let #7417 carry it along when it relocates the module.

@yuweih205
yuweih205 force-pushed the perf/flux2-fused-qk-norm-rope branch from 6d30c77 to 55d2328 Compare September 17, 2026 11:05
yuweih205 added a commit to yuweih205/vllm-omni that referenced this pull request Sep 17, 2026
… 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>
@yuweih205

Copy link
Copy Markdown
Contributor Author

@lishunyang12 Thanks again for the review. Your three points have been addressed: copying joint_attention_kwargs before updates, sharing _norm_rope_tile between both kernels, and skipping table packing when the fused path is unavailable.

Could you please re-review the current head (55d23287) and, if everything looks good, approve and help merge once the repository requirements are satisfied? All currently reported checks are green and the review threads are resolved. The shared op in this PR is also the prerequisite for the five follow-up model PRs.

Happy to address any remaining concerns. Thanks!

@hsliuustc0106 hsliuustc0106 added the Kernel optimization Codes related to optimize kernel execution to improve hardware utilization label Sep 20, 2026
yuweih205 added a commit to yuweih205/vllm-omni that referenced this pull request Sep 21, 2026
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>
yuweih205 added a commit to yuweih205/vllm-omni that referenced this pull request Sep 21, 2026
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>
yuweih205 added a commit to yuweih205/vllm-omni that referenced this pull request Sep 21, 2026
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>
yuweih205 added a commit to yuweih205/vllm-omni that referenced this pull request Sep 21, 2026
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>
yuweih205 added a commit to yuweih205/vllm-omni that referenced this pull request Sep 21, 2026
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>
@vllm-omni-review-bot

vllm-omni-review-bot commented Sep 21, 2026 •

Copy link
Copy Markdown

Omni ReviewBot triage note

Resolved as of 886b47c35f22: the high-priority or low-quality signal noted on an earlier commit no longer applies.

@yuweih205

Copy link
Copy Markdown
Contributor Author

The six-PR series is now synced with main (1b87115c): #7560 at 08dc20e6, followed by #7594, #7595, #7596, #7597 and #7600. GitHub currently reports all six as conflict-free. The dependency order is unchanged: this shared-op PR first, then the five independent model integrations.

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.

@yuweih205

yuweih205 commented Sep 21, 2026 •

Copy link
Copy Markdown
Contributor Author

Additional numerical/performance experiment (2026-09-21)

Keeping the current implementation at 08dc20e68; the bitwise-matching candidate below is experimental only and is not in this PR. This update adds evidence, not a code or gate change.

Numerical clarification. With vLLM 0.29.0's actual vllm_c RMSNorm provider, the residual Q/K differences are attributable to the different FP32 reduction order for sum(x²) (CUDA/CUB versus Triton's reduction tree). Floating-point addition is non-associative: these are deterministic, backend-dependent rounding differences, not run-to-run random noise. This does not imply that errors are unbiased or cannot accumulate through a model.

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 1e-6, Torch 2.13.0 / Triton 3.7.1 / vLLM 0.29.0 / CUDA 13.0. The reference is the original uncompiled vLLM RMSNorm → text-first cat (joint path) → actual RoPE chain, not another fused kernel. The shared source hash matches this head and #7595's shared operator.

  • In a controlled reduction-order probe, all 12 CUDA/CUB-order configurations had zero normalization/RoPE bit differences across 19 input cases; generic-reduction controls still differed.
  • An isolated copy of the real single/joint kernels reproducing that order then passed 42/42 strict Q/K/V checks: heads 24/48, batches 1/2, production/partial-tile shapes, packed-QKV strided views, multiple seeds, identity/random/signed weights and zero/small/large amplitudes. Integer-view comparisons include signed zero; all values were finite. Original/candidate repeatability checks also had zero differences.
  • Over the same matrix, the current fusion had 8,547 differing Q/K bit patterns out of 2,820,907,008 Q/K elements, maximum absolute difference 0.03125; V had zero differences. These are observed matrix results, not universal error bounds. The current PR is therefore not bitwise equal to the original chain.
  • A separate native-provider process also passed 42/42 with its own native-specific candidate; it is not used for the CUDA numbers below.

Performance tradeoff: three arms on identical inputs: original chain / current fusion / exact candidate. No profiler or torch.compile; 20 warmups, 100 invocations per sample, all six arm-order permutations. CUDA Graph measurements use 10 calls per replay and are normalized per call. All measured seeds and graph/eager outputs passed the candidate's exactness checks.

CUDA Graph replay, median ms per preparation call (not a full attention block or model):

B / text / image / heads Path Original Current fused Exact candidate Exact latency increase vs current
1 / 512 / 4096 / 24 joint 0.301403 0.055723 0.061743 10.8%
1 / 512 / 4096 / 24 single 0.186207 0.037914 0.050076 32.1%
1 / 512 / 4096 / 48 joint 0.581342 0.108935 0.119675 9.9%
1 / 512 / 4096 / 48 single 0.359613 0.074151 0.098641 33.0%
2 / 512 / 4096 / 24 joint 0.580660 0.109257 0.120144 10.0%
2 / 512 / 4096 / 24 single 0.359082 0.074519 0.099406 33.4%
1 / 512 / 16384 / 48 joint 2.118815 0.398219 0.434156 9.0%
1 / 512 / 16384 / 48 single 1.294909 0.273964 0.362015 32.1%

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: python benchmark_precision_r3.py --provider vllm_c --output-dir <fresh-experiment-output-dir> and a separate process with --provider native. The standalone harness, candidate copies and raw JSON are retained outside the PR; no benchmark scripts or experimental kernels were committed by this update.

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 vllm_c chain and original native chain against higher-precision references on the same H200/software stack. Here native means the PyTorch-based RMSNorm implementation running on the GPU, not a CPU baseline; vllm_c means vLLM's dedicated CUDA RMSNorm kernel. The tested default CUDA diffusion configuration prefers vllm_c; native is a separate control/fallback, not a synonym for the production baseline.

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:

Reference / stage Current fused Original CUDA Native
FP32 / RMSNorm 0.001614836425919 0.001614836425910 0.002073349678609
FP32 / final Q/K, full FP32 table 0.002625511161832 0.002625511219018 0.002929685189689
FP64 / RMSNorm 0.001614836802470 0.001614836802458 0.002073350372424
FP64 / final Q/K, full FP32 table 0.002625511451979 0.002625511509158 0.002929685732906
  • Versus native: current fusion lowers aggregate FP32-reference RMSE by 22.11% at RMSNorm and 10.38% at final Q/K. A separate reference holding the same rounded BF16 RoPE table values lowers final-Q/K RMSE by 13.12%. FP64 confirms the conclusion. Identity weights are essentially tied; the improvement is not an every-element or every-input guarantee.
  • Versus default CUDA: accuracy is essentially tied. Fusion is marginally worse at norm and marginally better after RoPE; the full-table Q/K RMSE reduction is only about 0.00000218%. Only 18,661 rotated values differ: fusion is closer to the FP32/full-table reference on 9,381, CUDA on 9,280, with the rest tied. No meaningful overall accuracy advantage over CUDA is established.

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: python precision_reference_r4.py --output-dir <fresh-experiment-output-dir> in the same fixed environment. The harness and raw per-case/aggregate JSON are retained outside the PR; no code, gate or PR-description changes accompany this result.

AI assistance: Codex prepared and ran the standalone experiments, checked the raw results, and drafted this update.

@vllm-omni-review-bot

Copy link
Copy Markdown

Omni ReviewBot: three questions on the performance claim

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

  • claim: It remained 3.58–4.88x faster than the original chain on those same workloads.

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.

@vllm-omni-review-bot

Copy link
Copy Markdown

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.

@vllm-omni-review-bot

Copy link
Copy Markdown
Omni ReviewBot routing record

Assigned Strict on zcode (GLM-5.3-Flash) under experiment fleet-strict-cursor-grok46-zcode-glm53flash-5050-c5-z10-20261002.

Copy link
Copy Markdown
Contributor Author

@vllm-omni-review-bot Short answers to the three questions:

  1. Bottleneck — the unfused preparation path is launch-heavy (RMSNorm -> cat -> RoPE, repeated across blocks). The per-stage/operator measurements in the PR show the fused prep path is 3.58–4.88x faster than the original chain under CUDA Graph replay on the measured H200 workloads.
  2. Value — on FLUX.2-dev, 1024², 28 steps, one H200, same-head A/B, p50 per image improved from 13.80 s to 13.42 s (-2.8%, about 13.6 ms/step). The optimization targets latency; no memory/cost claim is made.
  3. A/B / ablation — yes: same model/config/head, fusion gate off vs on, E-F-F-E order, 3 prompts × 2 seeds. Fused runs were bitwise-deterministic across repeats; eager-vs-fused PSNR was min 25.6 / median 35.8 / max 48.6 dB. The operator-level measurements isolate the prep fusion itself, while the e2e run shows its realized model-level effect.

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>
@yuweih205
yuweih205 force-pushed the perf/flux2-fused-qk-norm-rope branch from 08dc20e to 80d151c Compare October 8, 2026 04:15
Signed-off-by: HuangYuwei <yuweih205@gmail.com>
Signed-off-by: HuangYuwei <yuweih205@gmail.com>
@yuweih205
yuweih205 force-pushed the perf/flux2-fused-qk-norm-rope branch from 80d151c to 9ed72bf Compare October 8, 2026 06:13
@yuweih205

Copy link
Copy Markdown
Contributor Author

@lishunyang12 Your three review comments have been addressed on the current head (e702a525): joint_attention_kwargs is copied before updates, both kernels share _norm_rope_tile, and table packing is skipped before allocation when fusion is unavailable. All three review threads are resolved. Could you please re-review when you have a chance?

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.

@yuweih205

Copy link
Copy Markdown
Contributor Author

@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 1e-6): an experimental single/joint candidate passed 42/42 strict Q/K/V bitwise checks. Its preparation latency was 9.0–10.8% higher for joint and 32.1–33.4% higher for single than the current fusion. The candidate is not included in the current PRs; these are operator-level results on the documented configurations, so each model/provider would need its own validation.

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>
@yuweih205
yuweih205 force-pushed the perf/flux2-fused-qk-norm-rope branch from e702a52 to e0afca0 Compare October 8, 2026 07:42
@yuweih205

yuweih205 commented Oct 8, 2026 •

Copy link
Copy Markdown
Contributor Author

Supplemental eager bitwise measurements and queued production check

Existing 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 1e-6, seed 20261008, world/TP/SP size 1. The candidate has zero integer-view Q/K bit mismatches against the original eager CUDA RMSNorm/RoPE chain; routed V also matches exactly. Original and candidate repeat checks pass.

Path Original chain (ms) Current fusion (ms) Reference candidate (ms) Candidate saving vs original Candidate extra latency vs current
Joint 0.597912 0.111704 0.122705 79.48% 9.85%
Single 0.367308 0.075905 0.101037 72.49% 33.11%

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 VLLM_OMNI_FUSED_QK_NORM_ROPE_NUMERICS=fast to vllm_cuda_128; this resolves rms_norm_reduction="triton" → "vllm_cuda_128" and enable_fp_fusion=True → False. The common patch and validation records are prepared locally and have not been pushed to the six PRs. The retained table measured the isolated candidate.

Actual r3 target: prepared source 032950ed9533ae3565046defb2bf4d2f859938f9, using model_production_bitwise.py, SHA256 1d4eaff2c5de50b0c09a098adfab736bd936dba1c01653f870bd7dbb4b3efed9. The model's imported production functions are wrapped only for counting and receive unchanged arguments; no experimental operator is substituted.

The three eager arms use the same initialized model/state/inputs:

Arm MIN_TOKENS NUMERICS
Original chain 1000000000000 fast
Current fusion 0 fast
CUDA reference preset 0 vllm_cuda_128

H200, vLLM 0.29.0, Torch 2.13.0/CUDA 13.0, Triton 3.7.1; BF16, head dimension 128, B1/B2, seed 20261008, world/TP/SP size 1; two dual and two single blocks, 64 configured text tokens and a 16×16 image/video grid. Original RMSNorm is pinned to vllm_c; the harness records source/operator hashes and actual RMSNorm/RoPE forwards/providers. Complete output storage bits must match between the CUDA reference and original; every arm must repeat exactly and meet its fused-call counts before eager ABBA timing (20 warmup calls, 20 calls/sample, three rounds), including model RoPE generation and per-forward table packing.

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.

@yuweih205

Copy link
Copy Markdown
Contributor Author

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 032950ed9533ae3565046defb2bf4d2f859938f9. The numerical preset patch is still outside the public PR head; the exact shared patch is retained for review.

Settings: MIN_TOKENS=1000000000000, NUMERICS=fast for original; MIN_TOKENS=0, NUMERICS=vllm_cuda_128 for exact fusion, using the VLLM_OMNI_FUSED_QK_NORM_ROPE_ prefix. The preset selects rms_norm_reduction="vllm_cuda_128" and enable_fp_fusion=False. Model calls use the prepared operator directly.

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 20261008, world/TP/SP size 1. Two joint/dual blocks and two single blocks.
These timings include the complete shallow transformer forward, model RoPE generation and per-forward table packing. Text encoding, sampling and VAE are outside this measurement.

Batch Original chain (ms) Bitwise preset (ms) Forward latency reduction Output storage bits
1 38.450858 37.274188 3.060% 0 byte mismatches
2 76.559448 74.354437 2.880% 0 byte mismatches

The original RMSNorm provider is pinned and checked as vllm_c. Each arm repeats exactly; exact versus original is checked before and after timing. Source/operator hashes, state/input hashes, real norm/RoPE dispatch and fused-call counts are recorded. Provider hooks and counter wrappers are removed during timing.

Timing: H200 141 GB, vLLM 0.29.0, Torch 2.13.0+cu130, CUDA 13.0, Triton 3.7.1; CUDA-event medians, 10 warmup forwards per arm, three ABBA rounds, 10 calls/sample, six samples/arm. These are the stated random-weight eager cases; they do not establish pretrained full image/video generation parity.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

high priority high priority issue, needs to be done asap Kernel optimization Codes related to optimize kernel execution to improve hardware utilization

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants