Skip to content

Make SD3Transformer2DModel hidden states contiguous#14186

Open
menglcai wants to merge 2 commits into
huggingface:mainfrom
menglcai:menglcai/add_sd3_contiguous
Open

Make SD3Transformer2DModel hidden states contiguous#14186
menglcai wants to merge 2 commits into
huggingface:mainfrom
menglcai:menglcai/add_sd3_contiguous

Conversation

@menglcai

Copy link
Copy Markdown

What does this PR do?

Ensures hidden states are contiguous after each JointTransformerBlock call in SD3Transformer2DModel.

On ROCm, non-contiguous tensors produced by JointTransformerBlock cause performance degradation that accumulates across transformer blocks.

Benchmarks (SD3-medium, 512×512, 28 steps, fp16):

Platform PyTorch Baseline Patched Speedup
AMD Radeon(TM) 8060S Graphics (gfx1151) 2.12.0+rocm7.15.0a20260711 8503 ms 7470 ms +12.1%
AMD RX 9070 XT (gfx1201) 2.12.0+rocm7.15.0a20260713 2547 ms 2358 ms +7.4%
NVIDIA RTX 5090 2.12.1+cu132 1042 ms 1017 ms within noise

Before submitting

Who can review?

@sayakpaul @yiyixuxu

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@github-actions github-actions Bot added models size/S PR with diff < 50 LOC labels Jul 14, 2026
Comment on lines +322 to +324
hidden_states = hidden_states.contiguous()
if encoder_hidden_states is not None:
encoder_hidden_states = encoder_hidden_states.contiguous()

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Could also comment on the position in the block of code where contiguous() should be applied?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

You're right. After investigating, the non-contiguous tensor originates from PatchEmbed.forward - flatten(2).transpose(1, 2) (BCHW → BNC) produces a non-contiguous layout.

latent = latent.flatten(2).transpose(1, 2) # BCHW -> BNC

So the fix is more precisely placed right after self.pos_embed(), not inside the loop.

Updated the patch: moved .contiguous() to immediately after self.pos_embed(hidden_states), dropped the encoder_hidden_states call, and added a comment explaining the root cause.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

models size/S PR with diff < 50 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants