Skip to content

ZImageTransformer2D: Only build attention mask if seqlens are not equal#12955

Merged
yiyixuxu merged 1 commit intohuggingface:mainfrom
Beinsezii:beinsezii/zimg_optim
Mar 24, 2026
Merged

ZImageTransformer2D: Only build attention mask if seqlens are not equal#12955
yiyixuxu merged 1 commit intohuggingface:mainfrom
Beinsezii:beinsezii/zimg_optim

Conversation

@Beinsezii
Copy link
Copy Markdown
Contributor

What does this PR do?

Fixes a small performance regression for Z Image Turbo.

Basically just sets attn_mask to None when it would otherwise be all ones, which is always the case for Z Image Turbo where guidance_scale==1 for typical usage.

On an H100 this improves performance by about 4%, using AttentionBackendName._NATIVE_CUDNN.

Before submitting

Who can review?

@yiyixuxu or @sayakpaul probably

@sayakpaul
Copy link
Copy Markdown
Member

Cc: @JerryWu-code who contributed the model.

@Beinsezii
Copy link
Copy Markdown
Contributor Author

More relevant with new higher compute model

@Beinsezii
Copy link
Copy Markdown
Contributor Author

We've been running this in production for Turbo for a while

@Beinsezii Beinsezii force-pushed the beinsezii/zimg_optim branch from 02ae19d to b6d107a Compare January 28, 2026 00:25
@Beinsezii
Copy link
Copy Markdown
Contributor Author

👀

@dxqb
Copy link
Copy Markdown
Contributor

dxqb commented Mar 15, 2026

looks good to me
supercedes #13184

Copy link
Copy Markdown
Collaborator

@yiyixuxu yiyixuxu left a comment

Choose a reason for hiding this comment

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

thanks
sorry this slipped through

@HuggingFaceDocBuilderDev
Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@yiyixuxu yiyixuxu merged commit da6718f into huggingface:main Mar 24, 2026
11 checks passed
dxqb added a commit to Nerogar/OneTrainer that referenced this pull request Mar 24, 2026
speeds up Z-image by up to 25% because of huggingface/diffusers#12955

it speeds up attention, so the most benefit get higher resolutions. speed-up for 512px is not that much.
@Beinsezii Beinsezii deleted the beinsezii/zimg_optim branch March 24, 2026 22:11
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.

5 participants