Skip to content

ggml-cuda: generic DiT and VAE fusions for every family, plus a rope PDL race fix - #27

Merged
danielhanchen merged 10 commits into
masterfrom
perf/general-dit-vae
Oct 12, 2026
Merged

danielhanchen merged 10 commits into
masterfrom
perf/general-dit-vae

Conversation

@danielhanchen

@danielhanchen danielhanchen commented Oct 11, 2026 •

Copy link
Copy Markdown
Member

Problem

The CUDA levers from #17 to #25 were tuned on H3 and mostly sit in model-specific paths. Profiling the other Studio families on an RTX PRO 6000 Blackwell (sm120) and an L4 (sm89) showed where time still goes:

  • DiT glue: RoPE, the head-major relayout, the kv scale and the F16 cast of q / k / v run as separate ops before flash attention (Z-Image, FLUX, FLUX.2, Chroma, Qwen-Image). Modulation x*s + y chains add more passes.
  • KL VAE decode (SD 1.x, SDXL, FLUX, Z-Image, Chroma): GROUP_NORM runs one block per group and is then followed by separate MUL, ADD and SILU passes (GN plus elementwise is 25 to 35% of decode).
  • Wan-style VAE (Qwen-Image, Qwen-Image-2.1, Wan): every RMS norm is wrapped in two relayout copies (CONT is 31% of the Qwen-Image-2.1 decode).
  • rope_apply_f32 did not wait on its programmatic-dependent-launch predecessor. On sm90+ this made Z-Image renders differ between identical runs (PSNR 17 to 21 dB between two runs of the same seed on sm120). With GGML_CUDA_PDL=0 the output is stable, which pins the race.

Change

A new ggml patch scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch (applied after 0001 to 0006) plus small sd.cpp changes. Every fusion reproduces the stock rounding and reduction order, so outputs are bit-identical to the ops they replace. Each has a kill switch (=0).

change families switch
Rope::attention builds q / k / v with one rope_pe_permute op each (interleaved pairing mode added to the op, CPU and CUDA). Falls back to the old path for non-interleaved RoPE, no flash attention, sage, CPU backends, head dim below 64, or a backend that does not support the op FLUX.1, FLUX.2, Chroma, Z-Image, Qwen-Image, Wan and the other Rope::attention DiTs SD_FUSED_ROPE_ATTN
GROUP_NORM spread over the whole GPU for groups of 1024+ elements, fused with the affine MUL / ADD and an optional SILU SD 1.x / SDXL U-Net, all KL VAE decodes GGML_CUDA_GROUP_NORM_PAR
Channel RMS norm computed in place of CONT(PERMUTE) + RMS_NORM + MUL + PERMUTE + CONT, on layers large enough to win (16384+ positions for C < 1024, 65536+ for C >= 1024) Qwen-Image, Qwen-Image-2.1, Wan VAE GGML_CUDA_NORM_CHANNEL (=2 forces every size)
MUL + ADD (+ ADD) broadcast chains in one row-wise pass most DiTs GGML_CUDA_FUSE_MUL_ADD
ggml_cuda_pdl_sync() at the top of rope_apply_f32 every family using ggml rope on sm90+ none (correctness fix)
ggml_ext_conv_2d (im2col path) and Wan Conv2dBut3d skip a final cont that moves no data (batch 1 / single frame) all VAEs on the im2col path none

cmake/ggml.cmake detects ggml_rope_pe_permute_set_interleaved and defines SD_GGML_FUSED_ROPE_ATTN, so an unpatched ggml keeps the old graph. The switches are documented in docs/performance.md.

Results

Base = master (afac7a7) with patches 0001 to 0006, head = this branch, both built on the same VM in the same session. Studio defaults (--diffusion-fa --diffusion-conv-direct), seed 42, family default steps. Seconds, base runs then head runs, interleaved ABBA. Accuracy column compares decoded pixels of base and head.

RTX PRO 6000 Blackwell Server (sm120, Colab G4)

family size sampling base sampling head decode base decode head output vs base
Z-Image-Turbo Q4_K_M 1024 2.76 / 2.76 / 2.76 2.58 / 2.58 / 2.58 0.46 0.42 to 0.43 identical to base with GGML_CUDA_PDL=0; base itself is nondeterministic (PDL race), head is deterministic
Z-Image-Turbo Q4_K_M 512 1.03 / 1.03 / 1.03 0.98 / 0.99 / 0.98 0.27 0.26 as above
Z-Image-Turbo, --vae-tiling 7x7 1024 2.77 2.58 0.82 0.77 as above
FLUX.2-klein-4B Q4_K_M 1024 0.98 / 0.98 / 0.97 0.93 / 0.93 / 0.93 0.46 0.42 identical
FLUX.1-schnell Q4_K_M 1024 2.05 / 2.05 / 2.05 2.11 / 1.93 / 2.11 0.44 to 0.45 0.41 identical
SDXL base 1.0 1024 3.77 / 3.73 / 3.74 3.70 / 3.70 / 3.70 0.45 0.41 identical
SD 1.5 512 1.05 / 1.03 / 1.03 1.04 / 1.04 / 1.03 0.26 to 0.27 0.25 identical
Qwen-Image-2.1 Q4_K_M 1024 5.88 / 5.84 / 5.84 5.79 / 5.79 / 5.79 0.73 to 0.75 0.62 identical
Qwen-Image-2.1, --vae-tiling 7x7 1024 5.84 5.80 3.44 / 3.44 / 3.44 3.28 / 3.28 / 3.28 identical
Qwen-Image-2.1 Q4_K_M 1664x928 8.98 / 8.98 8.92 / 8.92 0.98 / 0.98 0.78 / 0.78 identical

L4 (sm89, Colab)

family size sampling base sampling head decode base decode head output vs base
Z-Image-Turbo Q4_K_M 1024 15.33 / 15.65 14.35 / 14.56 1.69 / 1.70 1.64 / 1.61 identical
Z-Image-Turbo Q4_K_M 512 4.41 / 4.42 4.13 / 4.13 0.57 / 0.58 0.56 / 0.56 identical
FLUX.1-schnell Q4_K_M 1024 11.15 / 11.02 10.61 / 10.48 1.58 / 1.63 1.52 / 1.54 identical
FLUX.2-klein-4B Q4_K_M 1024 4.89 / 4.92 4.85 / 4.70 1.68 / 1.70 1.60 / 1.60 identical
SDXL base 1.0 1024 13.26 / 13.50 13.32 / 13.45 1.58 / 1.62 1.52 / 1.53 identical
SD 1.5 512 2.74 / 2.73 2.73 / 2.75 0.52 / 0.51 0.49 / 0.49 identical
Qwen-Image-2.1 Q4_K_M 1024 33.05 32.59 to 32.67 3.29 to 3.36 3.22 identical

The L4 Qwen-Image-2.1 row was measured before the channel RMS norm fusion engaged on that VAE, so its decode gain is not included there.

Summary: DiT sampling 5 to 7% faster on Z-Image (both GPUs), 4 to 5% on FLUX.2-klein and FLUX.1 (FLUX.1 on sm120 is within noise), 1% on Qwen-Image-2.1, flat on the U-Net families. VAE decode 4 to 9% faster on KL VAEs, 15 to 20% on the Qwen-Image-2.1 VAE untiled and 5% tiled.

MiniMax-H3 and Wan 2.2, B200 (sm_100)

Base afac7a7 (patches 0001 to 0006) vs head (0001 to 0007), each applied to a clean ggml checkout, CI CUDA flags, exclusive card, interleaved ABBA, fresh process per run. H3: unsloth/MiniMax-H3-GGUF fl2va Q2_K with Studio's flags (--mode vid_gen --cfg-scale 1.0 --rng cpu --fps 24 --diffusion-fa), seed 11. Wan: Wan2.2-TI2V-5B Q4_K_M, euler, CFG 6, flow shift 3. Medians.

run output vs base sampling base / head VAE decode base / head total base / head
H3 640x384x24, 4 steps identical frames and audio 4.94 / 4.94 1.04 / 1.05 (tie) 7.98 / 7.95
H3 960x544x48, 4 steps, n=9 identical frames and audio 13.47 / 13.45 2.21 / 2.22 (tie) 17.82 / 17.89 (tie)
Wan2.2 5B 640x352x17, 8 steps identical to base with GGML_CUDA_PDL=0; base default differs run to run (8 of 8 distinct) 2.25 / 2.22 2.51 / 2.43 (-3.2%, CI -0.115 to -0.060 s) 5.84 / 5.82
  • H3 runs the new rope_pe_permute kernels and 25 fused MUL+ADD launches; Wan runs the interleaved fused rope attention and 55 channel RMS norm launches (nsys). With every switch off, both give the same output again.
  • Base is deterministic on H3 (its DiT does not use rope_apply) but not on Wan (two base runs about 25 dB apart); head is deterministic on both.
  • H3 VAE decode at 960x544 had a slow tail in some runs of both arms on this shared host (up to 3.38 s head, 2.34 s base in a follow-up set); the medians tie and it did not follow any switch.
  • Rechecked at the final head (with the overlap guard): Wan and H3 640 give the same hashes as above.

Tests

  • test-backend-ops on CUDA vs CPU (sm120), with local test cases for the new ops (not part of the patch): ROPE_PE_PERMUTE 120/120 (new interleaved cases), GROUP_NORM_AFFINE 16/16 (affine, SILU and reshape variants), RMS_NORM_CHANNEL 24/24 (C up to 2100, also with GGML_CUDA_NORM_CHANNEL=2 on sm100), MUL_ADD_CHAIN 12/12, GROUP_NORM, GROUP_NORM_MUL_ADD 2/2, FLASH_ATTN_EXT hsk=128 577/577.
  • Engagement: with GGML_CUDA_NORM_CHANNEL_DEBUG=1 all 36 channel RMS norms of the Qwen-Image-2.1 decode match the pattern and the 19 above the size gate fuse (untiled 1024); the kill-switch runs show the timing gains disappear with the switches off.
  • Decoded pixels identical base vs head for every row above, including tiled decodes (no seams introduced: same pixels as base).
  • With every switch set to 0 the head output is identical to base.
  • CPU-only build (GGML_NATIVE=OFF, the prebuilt CPU flags) of this branch compiles with all seven patches applied, and SD 1.5 at 256x256, 4 steps, seed 42 on CPU gives the same pixels as master (max abs diff 0).

Review fixes (0f53e72)

  • The fused RoPE + flash attention path now applies the --attn-scale override like the unfused path. Before, FLUX.2-klein with --attn-scale 0.0078125 gave the same pixels as no override on the fused path; now fused and unfused both give the unfused path's image, and without the override the image is unchanged.
  • The MUL+ADD(+ADD) and channel RMS norm matchers require GGML_TENSOR_FLAG_COMPUTE on every node they consume, as the GROUP_NORM matcher already did, so a graph built with ggml_build_forward_select never has a selected output skipped. sd.cpp does not build such graphs today.
  • The CPU interleaved ROPE_PE_PERMUTE evaluates the same expression as ROPE_APPLY, so FMA contraction cannot make them differ (6 of 12 F32 cases differed bitwise before on an -mfma build, 12 of 12 match now).
  • ggml_prec_set_acc replaces the deprecated precision setter.
  • A local test-fusion-select program (kept out of the patch, per CONTRIBUTING.md): selected-branch graphs on CPU and CUDA must match an expand-only graph bitwise and leave inactive outputs untouched, plus CPU RoPE parity. The previous head fails 16 of 40 on a B200, this head passes 40 of 40 (26 of 26 on a CPU-only build). test-backend-ops on the 0007 ops: 158 of 158 on both heads. FLUX.2-klein and SD1.5 renders are pixel-identical between the two heads, and 0001 to 0009 apply in order on a fresh ggml checkout.

Review fixes (f017bf3)

  • The GroupNorm affine / SiLU fusion now engages on real graphs. sd.cpp builds those as in-place MUL / ADD / SILU, which are views of the GroupNorm output, and the generic fusion check rejects views, so earlier heads only ran the parallel GroupNorm followed by separate elementwise kernels. The chain matcher now accepts a node that is a whole view of the GroupNorm tensor itself (offset 0, same data). On SD1.5 512 (U-Net plus VAE) the fused affine kernels go from 0 to 1494 launches, GroupNorm-chain launches drop from 18486 to 14389 and their device time from 85.2 to 66.4 ms (one nsys profile per arm). The decode timings above were measured before this fix, so they reflect the parallel GroupNorm alone; no new wall-clock number is claimed (shared card, too noisy).
  • Fused GroupNorm and channel RMS norm assert eps >= 0 like the unfused ops (a negative eps used to produce NaNs on the fused channel RMS path).
  • The channel RMS norm matcher requires both CONT_4D nodes to keep the input shape (a reshaping CONT used to fuse and give wrong values), and the GroupNorm chain only fuses when the output does not overlap the weight or bias and aliases the input only element for element.
  • Decoded pixels are identical between 122b715 and f017bf3 for FLUX.2-klein 512 and 1024 and SD1.5 512, and SD1.5 with GGML_CUDA_GROUP_NORM_PAR=0 matches too. The local test-fusion-select program goes from 46 of 48 to 48 of 48, and test-backend-ops on the 0007 ops plus GROUP_NORM passes 162 of 162 on both heads. 0001 to 0009 apply in order on a fresh ggml checkout.

Not covered

  • LTX, Wan 2.1 / A14B / I2V and the H3 Ref2VA, keyframe, offload and tiled-VAE modes were not run; H3 and Wan 2.2 TI2V-5B were run on a B200 only (above).
  • Chroma, Krea and Qwen-Image-2.1 Turbo were not timed separately; they share the FLUX / Qwen-Image paths above.
  • T4 (sm75) and A100 (sm80) timings: not run in this pass.
  • HIP and MUSA: the new fusions are switched off there. The interleaved rope_pe_permute path was neither built nor run for HIP, Vulkan or Metal; on those backends Rope::attention falls back when the op is unsupported.
  • The size gate of the channel RMS norm was tuned on B200 and checked on sm120; smaller GPUs reach the crossover earlier, so the gate is conservative there.

…parallel group norm, channel rms norm, mul+add)

Adds ggml patch 0007 on top of 0001-0006:
- rope_pe_permute gains an interleaved pairing mode (Z-Image, Flux, Chroma rope)
- rope_apply_f32 waits on the PDL predecessor (fixes nondeterministic output on sm90+)
- parallel GROUP_NORM for groups >= 1024 elements, fused with the affine MUL/ADD and SILU
- channel RMS norm fusion for CONT(PERMUTE) -> RMS_NORM -> MUL -> PERMUTE -> CONT (Wan/Qwen-Image VAE)
- MUL -> ADD [-> ADD] broadcast fusion for modulation
All fusions reproduce the stock rounding and reduction order, so outputs are bit-identical.
Each has a kill switch (GGML_CUDA_GROUP_NORM_PAR, GGML_CUDA_NORM_CHANNEL, GGML_CUDA_FUSE_MUL_ADD).
cmake detects the interleaved setter and defines SD_GGML_FUSED_ROPE_ATTN.
…conts

Rope::attention builds q/k/v through rope_pe_permute (interleaved mode) when flash attention is on
and the shapes match (Z-Image, Flux, Flux2, Chroma, Qwen-Image), falling back to the existing path otherwise.
ggml_ext_attention_prepared accepts an optional mask. SD_FUSED_ROPE_ATTN=0 disables it.
Conv2d im2col output and Wan Conv2dBut3d skip the final cont when the permute is a no-op (N == 1 / T == 1).
…chain

The fused kernel walks every channel in one thread per position, so it is slower than
CONT + RMS_NORM + MUL + CONT on small layers (VAE tiles). Fuse from 16384 positions when
C < 1024 and from 65536 when C >= 1024; GGML_CUDA_NORM_CHANNEL=2 forces it at every size.
Adds VAE-shaped perf cases for RMS_NORM_CHANNEL.
@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Oct 11, 2026 •

Copy link
Copy Markdown

Codex Review Summary

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

Review Status Commit Review trigger
📝 Code Review ✅ Completed 2026-10-12T05:44:21.861715Z 5da68dc New commits
ℹ️ About Codex in GitHub

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

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

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

@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Oct 11, 2026

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: f170601cf8

ℹ️ About Codex in GitHub

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

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

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment on lines +782 to +785
+ const char * x0 = (const char *) x->data;
+ const char * d0 = (const char *) cont2->data;
+ const bool partial_overlap = d0 != x0 && d0 < x0 + ggml_nbytes(x) && x0 < d0 + ggml_nbytes(cont2);
+ if (P > INT64_C(1) << 40 || partial_overlap) {

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Reject output buffers that overlap the weight

When w is produced inside the graph and the MUL is its final consumer, the allocator may reuse its buffer for cont2, because the unfused final CONT runs only after the MUL has finished reading w. This fusion instead launches at the first CONT and reads w while writing cont2; the overlap check covers only x, so such a reuse corrupts the normalization output. Use ggml_cuda_check_fusion_memory_ranges or explicitly reject overlap between cont2 and w before launching.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Fixed in aeb58f0: the channel RMS norm fusion now also rejects any overlap between its output (the second CONT) and the weight w, alongside the existing partial-overlap check against x, so a reused w buffer falls back to the unfused chain.

@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

1 similar comment
@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. 👍

Reviewed commit: aeb58f026f

ℹ️ About Codex in GitHub

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

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

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

…odes outside a selected graph branch

- Rope::attention's fused RoPE + flash path now applies ctx->attn_scale like ggml_ext_attention_ext,
  so --attn-scale keeps large K/V finite in F16 on both paths and they stay bit-identical
- ggml_ext_attention_prepared: ggml_prec_set_acc instead of the deprecated ggml_flash_attn_ext_set_prec
- 0007: the MUL+ADD(+ADD) and channel RMS norm matchers require GGML_TENSOR_FLAG_COMPUTE on every node
  they consume (graphs from ggml_build_forward_select), as ggml_can_fuse_ext does for GROUP_NORM
- 0007: CPU interleaved ROPE_PE_PERMUTE evaluates the same expression as ROPE_APPLY so FMA contraction
  cannot make them differ
- 0007: tests/test-fusion-select.cpp covers the selected-branch graphs and the CPU RoPE parity
@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

1 similar comment
@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 0f53e72072

ℹ️ About Codex in GitHub

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

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

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment on lines +1104 to +1105
+ set(TEST_TARGET test-fusion-select)
+ add_executable(${TEST_TARGET} ${TEST_TARGET}.cpp)

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P1 Badge Remove test code from the committed patch

This hunk adds test-fusion-select.cpp to the patched ggml test suite, and the same patch also extends test-backend-ops.cpp; however, the repository's CONTRIBUTING.md explicitly prohibits including test code or test scripts in commits or PRs and requires keeping them local. Remove these test-related patch hunks while retaining the verification results in the PR description.

AGENTS.md reference: AGENTS.md:L120-L122

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Removed in 3fb027e: 0007 no longer touches tests/ (the test-backend-ops cases and test-fusion-select are kept locally), and the results are in the PR description. 0001 to 0009 still apply in order on a fresh ggml checkout.

@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

1 similar comment
@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. Nice work!

Reviewed commit: 3fb027e9d7

ℹ️ About Codex in GitHub

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

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

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

- fused GroupNorm and channel RMS norm assert eps >= 0 like the unfused ops
- channel RMS norm: both CONTs must keep their PERMUTE's shape (a CONT_4D may reshape)
- GroupNorm affine: w and b may not overlap the fused output, x only element for element
  (the generic range check skips leaf sources)
- GroupNorm affine (+ SILU): accept the in-place chain sd.cpp builds (MUL / ADD / SILU are whole
  views of the GROUP_NORM result), keeping the COMPUTE, OUTPUT and use-count checks
@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. Another round soon, please!

Reviewed commit: f017bf3dda

ℹ️ About Codex in GitHub

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

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

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

@danielhanchen
danielhanchen merged commit 7612753 into master Oct 12, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant